| 13 | |
| 14 | class PPO(): |
| 15 | def __init__(self, policy, env, |
| 16 | learning_rate: float = 1e-5, |
| 17 | n_steps_total: int = 8192, |
| 18 | batch_size: int = 256, |
| 19 | n_epochs: int = 20, |
| 20 | gamma: float = 0.99, |
| 21 | gae_lambda: float = 0.9, |
| 22 | clip_range: float = 0.2, |
| 23 | clip_range_vf: float = None, |
| 24 | ent_coef: float = 0.05, |
| 25 | explore_coef: float = 0.05, |
| 26 | vf_coef: float = 0.5, |
| 27 | max_grad_norm: float = 0.5, |
| 28 | target_kl: float = 0.01, |
| 29 | update_adv=False, |
| 30 | lr_schedule_step=None, |
| 31 | start_num_timesteps: int = 0): |
| 32 | |
| 33 | self.policy = policy |
| 34 | self.env = env |
| 35 | self.learning_rate = learning_rate |
| 36 | self.n_steps_total = n_steps_total |
| 37 | self.n_steps = n_steps_total//env.num_envs |
| 38 | self.batch_size = batch_size |
| 39 | self.n_epochs = n_epochs |
| 40 | self.gamma = gamma |
| 41 | self.gae_lambda = gae_lambda |
| 42 | self.clip_range = clip_range |
| 43 | self.clip_range_vf = clip_range_vf |
| 44 | self.ent_coef = ent_coef |
| 45 | self.explore_coef = explore_coef |
| 46 | self.vf_coef = vf_coef |
| 47 | self.max_grad_norm = max_grad_norm |
| 48 | self.target_kl = target_kl |
| 49 | self.update_adv = update_adv |
| 50 | self.lr_schedule_step = lr_schedule_step |
| 51 | self.start_num_timesteps = start_num_timesteps |
| 52 | self.num_timesteps = start_num_timesteps |
| 53 | |
| 54 | self._last_obs = None |
| 55 | self._last_dones = None |
| 56 | self.ep_stat_buffer = None |
| 57 | |
| 58 | self.buffer = PpoBuffer(self.n_steps, self.env.observation_space, self.env.action_space, |
| 59 | gamma=self.gamma, gae_lambda=self.gae_lambda, n_envs=self.env.num_envs) |
| 60 | self.policy = self.policy.to(self.policy.device) |
| 61 | |
| 62 | model_parameters = filter(lambda p: p.requires_grad, self.policy.parameters()) |
| 63 | total_params = sum([np.prod(p.size()) for p in model_parameters]) |
| 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: |