| 9 | from .modules import BaseModule |
| 10 | |
| 11 | class PPOActor(nn.Module): |
| 12 | def __init__(self, |
| 13 | obs_dim_dict, |
| 14 | module_config_dict, |
| 15 | num_actions, |
| 16 | init_noise_std): |
| 17 | super(PPOActor, self).__init__() |
| 18 | |
| 19 | module_config_dict = self._process_module_config(module_config_dict, num_actions) |
| 20 | |
| 21 | self.actor_module = BaseModule(obs_dim_dict, module_config_dict) |
| 22 | |
| 23 | # Action noise |
| 24 | self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) |
| 25 | self.distribution = None |
| 26 | # disable args validation for speedup |
| 27 | Normal.set_default_validate_args = False |
| 28 | |
| 29 | def _process_module_config(self, module_config_dict, num_actions): |
| 30 | for idx, output_dim in enumerate(module_config_dict['output_dim']): |
| 31 | if output_dim == 'robot_action_dim': |
| 32 | module_config_dict['output_dim'][idx] = num_actions |
| 33 | return module_config_dict |
| 34 | |
| 35 | @property |
| 36 | def actor(self): |
| 37 | return self.actor_module |
| 38 | |
| 39 | @staticmethod |
| 40 | # not used at the moment |
| 41 | def init_weights(sequential, scales): |
| 42 | [torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in |
| 43 | enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))] |
| 44 | |
| 45 | def reset(self, dones=None): |
| 46 | pass |
| 47 | |
| 48 | def forward(self): |
| 49 | raise NotImplementedError |
| 50 | |
| 51 | @property |
| 52 | def action_mean(self): |
| 53 | return self.distribution.mean |
| 54 | |
| 55 | @property |
| 56 | def action_std(self): |
| 57 | return self.distribution.stddev |
| 58 | |
| 59 | @property |
| 60 | def entropy(self): |
| 61 | return self.distribution.entropy().sum(dim=-1) |
| 62 | |
| 63 | def update_distribution(self, actor_obs): |
| 64 | mean = self.actor(actor_obs) |
| 65 | self.distribution = Normal(mean, mean*0. + self.std) |
| 66 | |
| 67 | def act(self, actor_obs, **kwargs): |
| 68 | self.update_distribution(actor_obs) |
no outgoing calls
no test coverage detected