(self)
| 64 | self.action_dist.sample_weights(self.dist_sigma, batch_size=n_envs) |
| 65 | |
| 66 | def _build(self) -> None: |
| 67 | last_layer_dim_pi = self.features_extractor.features_dim |
| 68 | policy_net = [] |
| 69 | for layer_size in self.policy_head_arch: |
| 70 | policy_net.append(nn.Linear(last_layer_dim_pi, layer_size)) |
| 71 | policy_net.append(self.activation_fn()) |
| 72 | last_layer_dim_pi = layer_size |
| 73 | self.policy_head = nn.Sequential(*policy_net).to(self.device) |
| 74 | # mu->alpha/mean, sigma->beta/log_std (nn.Module, nn.Parameter) |
| 75 | self.dist_mu, self.dist_sigma = self.action_dist.proba_distribution_net(last_layer_dim_pi) |
| 76 | last_layer_dim_vf = self.features_extractor.features_dim |
| 77 | value_net = [] |
| 78 | for layer_size in self.value_head_arch: |
| 79 | value_net.append(nn.Linear(last_layer_dim_vf, layer_size)) |
| 80 | value_net.append(self.activation_fn()) |
| 81 | last_layer_dim_vf = layer_size |
| 82 | value_net.append(nn.Linear(last_layer_dim_vf, 1)) |
| 83 | self.value_head = nn.Sequential(*value_net).to(self.device) |
| 84 | # Init weights: use orthogonal initialization |
| 85 | # with small initial weight for the output |
| 86 | if self.ortho_init: |
| 87 | # TODO: check for features_extractor |
| 88 | # Values from stable-baselines. |
| 89 | # feature_extractor/mlp values are |
| 90 | # originally from openai/baselines (default gains/init_scales). |
| 91 | module_gains = { |
| 92 | # self.features_extractor: np.sqrt(2), |
| 93 | self.policy_head: np.sqrt(2), |
| 94 | self.value_head: np.sqrt(2) |
| 95 | # self.action_net: 0.01, |
| 96 | } |
| 97 | for module, gain in module_gains.items(): |
| 98 | module.apply(partial(self.init_weights, gain=gain)) |
| 99 | |
| 100 | self.optimizer = self.optimizer_class(self.parameters(), **self.optimizer_kwargs) |
| 101 | |
| 102 | def _get_features(self, birdview: th.Tensor, state: th.Tensor) -> th.Tensor: |
| 103 | """ |
no test coverage detected