| 230 | return th.mean(entropy_loss) |
| 231 | |
| 232 | def exploration_loss(self, exploration_suggests) -> th.Tensor: |
| 233 | # [('stop'/'go'/None, 'turn'/'straight'/None)] |
| 234 | # (batch_size, action_dim) |
| 235 | alpha = self.distribution.concentration1.detach().clone() |
| 236 | beta = self.distribution.concentration0.detach().clone() |
| 237 | |
| 238 | for i, (acc_suggest, steer_suggest) in enumerate(exploration_suggests): |
| 239 | if acc_suggest != '': |
| 240 | beta[i, 0] = self.acc_exploration_dist[acc_suggest][0] |
| 241 | alpha[i, 0] = self.acc_exploration_dist[acc_suggest][1] |
| 242 | if steer_suggest != '': |
| 243 | beta[i, 1] = self.steer_exploration_dist[steer_suggest][0] |
| 244 | alpha[i, 1] = self.steer_exploration_dist[steer_suggest][1] |
| 245 | |
| 246 | dist_ent = Beta(alpha, beta) |
| 247 | |
| 248 | exploration_loss = th.distributions.kl_divergence(self.distribution, dist_ent) |
| 249 | return th.mean(exploration_loss) |
| 250 | |
| 251 | def sample(self) -> th.Tensor: |
| 252 | # Reparametrization trick to pass gradients |