| 130 | |
| 131 | |
| 132 | class Stage2ExportWrapper(torch.nn.Module): |
| 133 | def __init__(self, model: ActorCritic, rms_obs: RunningMeanStd, rms_hist: RunningMeanStd): |
| 134 | super().__init__() |
| 135 | self.model = model |
| 136 | self.rms_obs = rms_obs |
| 137 | self.rms_hist = rms_hist |
| 138 | |
| 139 | def forward(self, obs: torch.Tensor, proprio_hist: torch.Tensor) -> torch.Tensor: |
| 140 | obs_n = self.rms_obs(obs) |
| 141 | hist_n = self.rms_hist(proprio_hist) |
| 142 | mu = self.model.act_inference({"obs": obs_n, "proprio_hist": hist_n}) |
| 143 | return torch.clamp(mu, -1.0, 1.0) |
| 144 | |
| 145 | |
| 146 | class Stage1ExportWrapper(torch.nn.Module): |