(self)
| 27 | self._load_weight() |
| 28 | |
| 29 | def _load_weight(self): |
| 30 | rl_state_dict = torch.load(self.config.rl_ckpt, map_location='cpu')['policy_state_dict'] |
| 31 | self._load_state_dict(self.model.value_branch_traj, rl_state_dict, 'value_head') |
| 32 | self._load_state_dict(self.model.value_branch_ctrl, rl_state_dict, 'value_head') |
| 33 | self._load_state_dict(self.model.dist_mu, rl_state_dict, 'dist_mu') |
| 34 | self._load_state_dict(self.model.dist_sigma, rl_state_dict, 'dist_sigma') |
| 35 | |
| 36 | def _load_state_dict(self, il_net, rl_state_dict, key_word): |
| 37 | rl_keys = [k for k in rl_state_dict.keys() if key_word in k] |
no test coverage detected