| 99 | actor = copy.deepcopy(inference_model["actor"]).to("cpu") |
| 100 | |
| 101 | class PPOEncoderWrapper(nn.Module): |
| 102 | def __init__(self, actor): |
| 103 | """ |
| 104 | model: The original PyTorch model. |
| 105 | input_keys: List of input names as keys for the input dictionary. |
| 106 | """ |
| 107 | super(PPOEncoderWrapper, self).__init__() |
| 108 | self.actor = actor |
| 109 | |
| 110 | def forward(self, inputs): |
| 111 | """ |
| 112 | Dynamically creates a dictionary from the input keys and args. |
| 113 | """ |
| 114 | actor_obs, future_motion_targets, history = inputs |
| 115 | motion_embedding = self.actor.motion_encoder(future_motion_targets) |
| 116 | latent = self.actor.history_encoder(history) |
| 117 | input_for_actor = torch.cat([actor_obs, motion_embedding, latent], dim=-1) |
| 118 | return self.actor.actor_module(input_for_actor) |
| 119 | |
| 120 | wrapper = PPOEncoderWrapper(actor) |
| 121 | example_input_list = [ |
no outgoing calls
no test coverage detected