(self, env: VecEnv, callback: BaseCallback,
rollout_buffer: PpoBuffer, n_rollout_steps: int)
| 64 | print(f'trainable parameters: {total_params/1000000:.2f}M') |
| 65 | |
| 66 | def collect_rollouts(self, env: VecEnv, callback: BaseCallback, |
| 67 | rollout_buffer: PpoBuffer, n_rollout_steps: int) -> bool: |
| 68 | assert self._last_obs is not None, "No previous observation was provided" |
| 69 | n_steps = 0 |
| 70 | rollout_buffer.reset() |
| 71 | |
| 72 | self.action_statistics = [] |
| 73 | self.mu_statistics = [] |
| 74 | self.sigma_statistics = [] |
| 75 | |
| 76 | while n_steps < n_rollout_steps: |
| 77 | actions, values, log_probs, mu, sigma, _ = self.policy.forward(self._last_obs) |
| 78 | self.action_statistics.append(actions) |
| 79 | self.mu_statistics.append(mu) |
| 80 | self.sigma_statistics.append(sigma) |
| 81 | |
| 82 | new_obs, rewards, dones, infos = env.step(actions) |
| 83 | |
| 84 | if callback.on_step() is False: |
| 85 | return False |
| 86 | |
| 87 | # update_info_buffer |
| 88 | for i in np.where(dones)[0]: |
| 89 | self.ep_stat_buffer.append(infos[i]['episode_stat']) |
| 90 | |
| 91 | n_steps += 1 |
| 92 | self.num_timesteps += env.num_envs |
| 93 | |
| 94 | rollout_buffer.add(self._last_obs, actions, rewards, self._last_dones, values, log_probs, mu, sigma, infos) |
| 95 | self._last_obs = new_obs |
| 96 | self._last_dones = dones |
| 97 | |
| 98 | last_values = self.policy.forward_value(self._last_obs) |
| 99 | rollout_buffer.compute_returns_and_advantage(last_values, dones=self._last_dones) |
| 100 | |
| 101 | return True |
| 102 | |
| 103 | def train(self): |
| 104 | for param_group in self.policy.optimizer.param_groups: |
no test coverage detected