| 79 | raise NotImplementedError |
| 80 | |
| 81 | def initial_inference(self, obs) -> NetworkOutput: |
| 82 | num = obs.size(0) |
| 83 | |
| 84 | state = self.representation(obs) |
| 85 | actor_logit, value = self.prediction(state) |
| 86 | |
| 87 | if not self.training: |
| 88 | # if not in training, obtain the scalars of the value/reward |
| 89 | value = self.inverse_value_transform(value).detach().cpu().numpy() |
| 90 | state = state.detach().cpu().numpy() |
| 91 | actor_logit = actor_logit.detach().cpu().numpy() |
| 92 | # zero initialization for reward (value prefix) hidden states |
| 93 | reward_hidden = (torch.zeros(1, num, self.lstm_hidden_size).detach().cpu().numpy(), |
| 94 | torch.zeros(1, num, self.lstm_hidden_size).detach().cpu().numpy()) |
| 95 | else: |
| 96 | # zero initialization for reward (value prefix) hidden states |
| 97 | reward_hidden = (torch.zeros(1, num, self.lstm_hidden_size).to('cuda'), torch.zeros(1, num, self.lstm_hidden_size).to('cuda')) |
| 98 | |
| 99 | return NetworkOutput(value, [0. for _ in range(num)], actor_logit, state, reward_hidden) |
| 100 | |
| 101 | def recurrent_inference(self, hidden_state, reward_hidden, action) -> NetworkOutput: |
| 102 | state, reward_hidden, value_prefix = self.dynamics(hidden_state, reward_hidden, action) |