| 144 | |
| 145 | |
| 146 | class Stage1ExportWrapper(torch.nn.Module): |
| 147 | def __init__(self, model: ActorCritic, rms_obs: RunningMeanStd): |
| 148 | super().__init__() |
| 149 | self.model = model |
| 150 | self.rms_obs = rms_obs |
| 151 | |
| 152 | def forward(self, obs: torch.Tensor, priv_info: torch.Tensor) -> torch.Tensor: |
| 153 | obs_n = self.rms_obs(obs) |
| 154 | mu = self.model.act_inference({"obs": obs_n, "priv_info": priv_info}) |
| 155 | return torch.clamp(mu, -1.0, 1.0) |
| 156 | |
| 157 | |
| 158 | def _save_deploy_meta( |