| 199 | self.device = 'cpu' |
| 200 | |
| 201 | def proba_distribution_net(self, latent_dim: int) -> Tuple[nn.Module, nn.Module]: |
| 202 | |
| 203 | linear_alpha = nn.Linear(latent_dim, self.action_dim) |
| 204 | linear_beta = nn.Linear(latent_dim, self.action_dim) |
| 205 | |
| 206 | if self.dist_init is not None: |
| 207 | # linear_alpha.weight.data.fill_(0.01) |
| 208 | # linear_beta.weight.data.fill_(0.01) |
| 209 | # acc |
| 210 | linear_alpha.bias.data[0] = self.dist_init[0][1] |
| 211 | linear_beta.bias.data[0] = self.dist_init[0][0] |
| 212 | # steer |
| 213 | linear_alpha.bias.data[1] = self.dist_init[1][1] |
| 214 | linear_beta.bias.data[1] = self.dist_init[1][0] |
| 215 | |
| 216 | alpha = nn.Sequential(linear_alpha, nn.Softplus()) |
| 217 | beta = nn.Sequential(linear_beta, nn.Softplus()) |
| 218 | return alpha, beta |
| 219 | |
| 220 | def proba_distribution(self, alpha, beta): |
| 221 | self.distribution = Beta(alpha, beta) |