PPO v2

import%20marimo%0A%0A__generated_with%20%3D%20%220.25.2%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20base64%0A%20%20%20%20import%20io%0A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20torch.distributions%20as%20D%0A%20%20%20%20import%20gymnasium%20as%20gym%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20from%20PIL%20import%20Image%0A%0A%20%20%20%20return%20D%2C%20Image%2C%20base64%2C%20gym%2C%20io%2C%20mo%2C%20nn%2C%20np%2C%20plt%2C%20torch%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Solving%20Continuous%20Control%20with%20PPO%0A%0A%20%20%20%20We%20aim%20to%20solve%20continuous-action%20%60gymnasium%60%20environments%20using%20Proximal%20Policy%20Optimisation.%0A%0A%20%20%20%20A2C%20threw%20each%20experience%20away%20immediately%20after%20using%20it%20for%20a%20single%20gradient%20step.%0A%20%20%20%20PPO%20instead%20collects%20a%20large%20batch%20of%20experience%20and%20reuses%20it%20for%20several%20passes%20of%20gradient%20descent.%0A%20%20%20%20The%20difficulty%20is%20that%20after%20the%20first%20pass%20the%20data%20no%20longer%20comes%20from%20the%20current%20policy%2C%20so%20PPO%20clips%20the%20update%20to%20stop%20the%20policy%20drifting%20too%20far%20from%20the%20one%20that%20gathered%20the%20data.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Set%20up%20the%20Actor%20and%20Critic%20networks.%0A%20%20%20%20Both%20take%20the%20state%20as%20input%20and%20have%202%20hidden%20layers%20with%20%24%5Ctanh%24%20activations.%0A%0A%20%20%20%20The%20action%20space%20is%20now%20continuous%2C%20so%20the%20actor%20cannot%20output%20logits%20over%20a%20finite%20set%20of%20actions%20as%20it%20did%20in%20Acrobot.%0A%20%20%20%20Instead%20it%20outputs%20a%20mean%20%24%5Cmu(s)%24%20for%20each%20action%20dimension%20and%20we%20sample%20%24a%20%5Csim%20%5Cmathcal%7BN%7D(%5Cmu(s)%2C%20%5Csigma%5E2)%24.%0A%20%20%20%20The%20standard%20deviation%20%24%5Csigma%24%20is%20a%20learnable%20parameter%20in%20its%20own%20right%20and%20does%20not%20depend%20on%20the%20state%2C%20so%20the%20agent%20can%20steadily%20reduce%20its%20own%20exploration%20as%20it%20becomes%20more%20confident.%0A%20%20%20%20We%20store%20%24%5Cln%20%5Csigma%24%20rather%20than%20%24%5Csigma%24%20so%20the%20parameter%20is%20unconstrained%20while%20%24%5Csigma%20%3D%20e%5E%7B%5Cln%20%5Csigma%7D%24%20stays%20positive.%0A%0A%20%20%20%20%60act()%60%20is%20used%20while%20collecting%20data%2C%20and%20%60evaluate()%60%20recomputes%20%24%5Cln%20%5Cpi_%5Ctheta(a%20%5Cvert%7B%7D%20s)%24%20for%20actions%20that%20have%20already%20been%20taken%2C%20this%20time%20under%20the%20updated%20policy.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(D%2C%20nn%2C%20np%2C%20torch)%3A%0A%20%20%20%20def%20alter_layer(layer%2C%20orth_scaling%3Dnp.sqrt(2)%2C%20bias_init%3D0)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Force%20the%20weights%20to%20form%20an%20orthogonal%20matrix%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Prevents%20the%20scaling%20of%20signals%20down%20to%20zero%20or%20up%20to%20infinity%20as%20it%20passes%20through%20the%20layers%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Scaled%20by%20sqrt(2)%20to%20account%20for%20the%20tanh%20activations%20(we%20didn't%20need%20this%20in%20A2C%20because%20the%20activations%20were%20linear)%0A%20%20%20%20%20%20%20%20%20%20%20%20nn.init.orthogonal_(layer.weight%2C%20orth_scaling)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Force%20the%20bias%20tensors%20to%20be%20a%20constant%20prescribed%20by%20%60bias_init%60%0A%20%20%20%20%20%20%20%20%20%20%20%20nn.init.constant_(layer.bias%2C%20bias_init)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20layer%0A%0A%20%20%20%20class%20ActorCritic(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20state_dim%2C%20action_dim%2C%20neurons_per_layer%2C%20sigma_init)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.actor%20%3D%20nn.Sequential(%20%23%20Automatically%20connects%20the%20layers%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(state_dim%2C%20neurons_per_layer))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nn.Tanh()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(neurons_per_layer%2C%20neurons_per_layer))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nn.Tanh()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Outputs%20a%20mean%20for%20each%20action%20dimension%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Scaling%20the%20output%20by%200.01%20keeps%20the%20mean%20action%20at%20the%20beginning%20close%20to%20zero%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Exploration%20is%20encouraged%20organically%20without%20leaning%20towards%20one%20specific%20side%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(neurons_per_layer%2C%20action_dim)%2C%20orth_scaling%3D0.01)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.critic%20%3D%20nn.Sequential(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(state_dim%2C%20neurons_per_layer))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nn.Tanh()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(neurons_per_layer%2C%20neurons_per_layer))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nn.Tanh()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Only%20output%20V(s)%2C%20scale%20by%201%20because%20no%20activation%20following%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alter_layer(nn.Linear(neurons_per_layer%2C%201)%2C%20orth_scaling%3D1)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Separate%20learnable%20parameter%20of%20log(sigma)%20for%20each%20action%0A%20%20%20%20%20%20%20%20%20%20%20%20log_sigma_init%20%3D%20np.log(sigma_init)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Populate%20a%20tensor%20with%20an%20entry%20for%20each%20action_dim%20with%20log_sigma_init%0A%20%20%20%20%20%20%20%20%20%20%20%20log_sigma_init%20%3D%20torch.full((action_dim%2C)%2C%20log_sigma_init%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.log_sigma%20%3D%20nn.Parameter(log_sigma_init)%0A%0A%20%20%20%20%20%20%20%20%23%20Gets%20a%20distribution%20representing%20the%20actor's%20predicted%20best%20action%20over%20the%20action%20space%0A%20%20%20%20%20%20%20%20def%20dist(self%2C%20state)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20mean%20%3D%20self.actor(state)%0A%20%20%20%20%20%20%20%20%20%20%20%20sigma%20%3D%20torch.exp(self.log_sigma)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20D.Normal(mean%2C%20sigma)%0A%0A%20%20%20%20%20%20%20%20def%20value(self%2C%20state)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.critic(state).squeeze(-1)%0A%0A%20%20%20%20%20%20%20%20%40torch.no_grad()%0A%20%20%20%20%20%20%20%20def%20act(self%2C%20state)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20dist%20%3D%20self.dist(state)%0A%20%20%20%20%20%20%20%20%20%20%20%20action%20%3D%20dist.sample()%0A%20%20%20%20%20%20%20%20%20%20%20%20log_prob%20%3D%20dist.log_prob(action).sum(dim%3D-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20value_estimate%20%3D%20self.value(state)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20action%2C%20log_prob%2C%20value_estimate%0A%0A%20%20%20%20%20%20%20%20def%20evaluate(self%2C%20state%2C%20action)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20dist%20%3D%20self.dist(state)%0A%20%20%20%20%20%20%20%20%20%20%20%20log_prob%20%3D%20dist.log_prob(action).sum(dim%3D-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20value_estimate%20%3D%20self.value(state)%0A%20%20%20%20%20%20%20%20%20%20%20%20entropy%20%3D%20dist.entropy().sum(dim%3D-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20log_prob%2C%20entropy%2C%20value_estimate%0A%0A%20%20%20%20%20%20%20%20%40torch.no_grad()%0A%20%20%20%20%20%20%20%20def%20greedy_action(self%2C%20state)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.actor(state)%0A%0A%20%20%20%20return%20(ActorCritic%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20buffer%20stores%20one%20rollout%20of%20%60buffer_size%60%20transitions%20and%20then%20converts%20them%20into%20training%20targets.%0A%0A%20%20%20%20Advantages%20are%20computed%20using%20Generalised%20Advantage%20Estimation%2C%20an%20exponentially%20weighted%20average%20of%20every%20%24n%24-step%20advantage%20estimate%2C%20evaluated%20by%20a%20single%20backwards%20pass%20through%20the%20rollout%3A%0A%0A%20%20%20%20%24%24%5Cdelta_t%20%3D%20r_t%20%2B%20%5Cgamma%20V(s_%7Bt%2B1%7D)(1-d_t)%20-%20V(s_t)%2C%20%5Cqquad%20A_t%20%3D%20%5Cdelta_t%20%2B%20%5Cgamma%20%5Clambda%20(1-d_t)%20A_%7Bt%2B1%7D%24%24%0A%0A%20%20%20%20Setting%20%24%5Clambda%20%3D%200%24%20recovers%20the%20one-step%20A2C%20advantage%20(low%20variance%2C%20high%20bias)%2C%20while%20%24%5Clambda%20%3D%201%24%20gives%20the%20full%20Monte-Carlo%20return%20(high%20variance%2C%20low%20bias).%0A%20%20%20%20The%20critic's%20target%20is%20then%20simply%20%24A_t%20%2B%20V(s_t)%24.%0A%0A%20%20%20%20%60get_minibatch()%60%20normalises%20the%20advantages%20to%20zero%20mean%20and%20unit%20variance%2C%20which%20keeps%20the%20size%20of%20the%20policy%20update%20independent%20of%20the%20scale%20of%20the%20environment's%20rewards.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(torch)%3A%0A%20%20%20%20class%20RolloutBuffer%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20buffer_size%2C%20state_dim%2C%20action_dim%2C%20gamma%2C%20lam%2C%20minibatch_size)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Generated%20in%20the%20data%20collection%20phase%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer_size%20%3D%20buffer_size%0A%20%20%20%20%20%20%20%20%20%20%20%20self.step%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20self.states%20%3D%20torch.zeros((buffer_size%2C%20state_dim)%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.actions%20%3D%20torch.zeros((buffer_size%2C%20action_dim)%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.log_probs%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Store%20booleans%20as%20floats%20to%20prevent%20messy%20casts%20later%20on%0A%20%20%20%20%20%20%20%20%20%20%20%20self.dones%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rewards%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.values%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Calculated%20after%20data%20collection%0A%20%20%20%20%20%20%20%20%20%20%20%20self.advantages%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.returns%20%3D%20torch.zeros(buffer_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.gamma%20%3D%20gamma%0A%20%20%20%20%20%20%20%20%20%20%20%20self.lam%20%3D%20lam%0A%20%20%20%20%20%20%20%20%20%20%20%20self.minibatch_size%20%3D%20minibatch_size%0A%0A%20%20%20%20%20%20%20%20def%20insert(self%2C%20state%2C%20action%2C%20log_prob%2C%20terminated%2C%20truncated%2C%20reward%2C%20value)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.states%5Bself.step%5D%20%3D%20state%0A%20%20%20%20%20%20%20%20%20%20%20%20self.actions%5Bself.step%5D%20%3D%20action%0A%20%20%20%20%20%20%20%20%20%20%20%20self.log_probs%5Bself.step%5D%20%3D%20log_prob%0A%20%20%20%20%20%20%20%20%20%20%20%20self.dones%5Bself.step%5D%20%3D%20truncated%20or%20terminated%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rewards%5Bself.step%5D%20%3D%20reward%0A%20%20%20%20%20%20%20%20%20%20%20%20self.values%5Bself.step%5D%20%3D%20value%0A%20%20%20%20%20%20%20%20%20%20%20%20self.step%20%3D%20(self.step%20%2B%201)%20%25%20self.buffer_size%0A%0A%20%20%20%20%20%20%20%20def%20compute_gae(self%2C%20last_value)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20next_advantage%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20next_value%20%3D%20last_value%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20t%20in%20reversed(range(self.buffer_size))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20It%20is%20okay%20to%20use%20done%20here%20instead%20of%20terminated%20because%20it%20prevents%20double%20bootstrapping%20from%20the%20update_cycle%20method%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20delta%20%3D%20self.rewards%5Bt%5D%20%2B%20self.gamma%20*%20next_value%20*%20(1%20-%20self.dones%5Bt%5D)%20-%20self.values%5Bt%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.advantages%5Bt%5D%20%3D%20delta%20%2B%20self.gamma%20*%20self.lam%20*%20(1%20-%20self.dones%5Bt%5D)%20*%20next_advantage%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.returns%5Bt%5D%20%3D%20self.advantages%5Bt%5D%20%2B%20self.values%5Bt%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20next_advantage%20%3D%20self.advantages%5Bt%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20next_value%20%3D%20self.values%5Bt%5D%0A%0A%20%20%20%20%20%20%20%20def%20new_epoch(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.batch_indices%20%3D%20torch.randperm(self.buffer_size)%0A%0A%20%20%20%20%20%20%20%20def%20minibatches_per_epoch(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20self.buffer_size%20%25%20self.minibatch_size%20%3D%3D%200%2C%20%22Minibatch%20size%20must%20divide%20buffer%20size%22%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.buffer_size%20%2F%2F%20self.minibatch_size%0A%0A%20%20%20%20%20%20%20%20def%20get_minibatch(self%2C%20minibatch_idx)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20start_index%20%3D%20minibatch_idx%20*%20self.minibatch_size%0A%20%20%20%20%20%20%20%20%20%20%20%20end_index%20%3D%20(minibatch_idx%20%2B%201)%20*%20self.minibatch_size%0A%20%20%20%20%20%20%20%20%20%20%20%20indices%20%3D%20self.batch_indices%5Bstart_index%3Aend_index%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Normalise%20the%20advantages%0A%20%20%20%20%20%20%20%20%20%20%20%20adv%20%3D%20self.advantages%5Bindices%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20adv%20%3D%20(adv%20-%20adv.mean())%20%2F%20(adv.std()%20%2B%201e-8)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.states%5Bindices%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.actions%5Bindices%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.log_probs%5Bindices%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20adv%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.returns%5Bindices%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.values%5Bindices%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20return%20(RolloutBuffer%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Collect%20a%20rollout%20by%20running%20the%20current%20policy%20for%20exactly%20%60buffer_size%60%20steps.%0A%0A%20%20%20%20Because%20a%20rollout%20is%20a%20fixed%20number%20of%20steps%20rather%20than%20a%20fixed%20number%20of%20episodes%2C%20an%20episode%20can%20straddle%20the%20boundary%20between%20two%20rollouts.%0A%20%20%20%20We%20therefore%20thread%20the%20current%20state%20and%20the%20running%20episode%20total%20through%20to%20the%20next%20call.%0A%0A%20%20%20%20Termination%20and%20truncation%20must%20be%20treated%20differently.%0A%20%20%20%20A%20terminated%20episode%20genuinely%20has%20no%20future%20reward%2C%20but%20a%20truncated%20one%20was%20cut%20off%20by%20a%20time%20limit%20and%20would%20have%20carried%20on%2C%20so%20we%20add%20the%20bootstrapped%20%24%5Cgamma%20V(s')%24%20to%20its%20final%20reward%20instead%20of%20pretending%20the%20future%20is%20worth%20nothing.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(ActorCritic%2C%20RolloutBuffer%2C%20gym%2C%20torch)%3A%0A%20%20%20%20%23%20Data%20collection%20phase%0A%20%20%20%20%40torch.no_grad()%20%23%20Make%20sure%20these%20calculations%20do%20not%20affect%20the%20gradients%0A%20%20%20%20def%20collect_rollout(agent%20%3A%20ActorCritic%2C%20buffer%20%3A%20RolloutBuffer%2C%20env%20%3A%20gym.Env%2C%20initial_state%2C%20initial_episode_reward)%3A%0A%20%20%20%20%20%20%20%20cur_state%20%3D%20initial_state%0A%20%20%20%20%20%20%20%20%23%20Carried%20in%20from%20the%20previous%20rollout%20because%20an%20episode%20can%20straddle%20the%20boundary%20between%20two%20rollouts%0A%20%20%20%20%20%20%20%20total_reward%20%3D%20initial_episode_reward%0A%20%20%20%20%20%20%20%20%23%20Every%20episode%20that%20actually%20finished%20inside%20this%20rollout%0A%20%20%20%20%20%20%20%20episode_rewards%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20for%20_%20in%20range(buffer.buffer_size)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20action%2C%20log_prob%2C%20value_estimate%20%3D%20agent.act(cur_state)%0A%20%20%20%20%20%20%20%20%20%20%20%20next_state%2C%20reward%2C%20terminated%2C%20truncated%2C%20_%20%3D%20env.step(action.numpy())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Accumulate%20the%20raw%20environment%20reward%20before%20the%20truncation%20bootstrap%20is%20added%20below%0A%20%20%20%20%20%20%20%20%20%20%20%20total_reward%20%2B%3D%20reward%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Fix%20the%20data-types%0A%20%20%20%20%20%20%20%20%20%20%20%20next_state%20%3D%20torch.tensor(next_state%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20reward%20%3D%20torch.tensor(reward%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20done%20%3D%20terminated%20or%20truncated%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20truncated%3A%20%23%20Do%20not%20punish%20the%20agent%20for%20being%20truncated%20rather%20than%20terminated%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20reward%20%2B%3D%20buffer.gamma%20*%20agent.value(next_state)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Add%20to%20the%20buffer%0A%20%20%20%20%20%20%20%20%20%20%20%20buffer.insert(cur_state%2C%20action%2C%20log_prob%2C%20terminated%2C%20truncated%2C%20reward%2C%20value_estimate)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20done%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20The%20episode%20is%20over%2C%20so%20bank%20its%20total%20and%20start%20counting%20again%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20episode_rewards.append(total_reward)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20total_reward%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cur_state%2C%20_%20%3D%20env.reset()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cur_state%20%3D%20torch.tensor(cur_state%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cur_state%20%3D%20next_state%0A%0A%20%20%20%20%20%20%20%20%23%20Get%20the%20value%20estimate%20for%20the%20state%20after%20the%20last%20experience%0A%20%20%20%20%20%20%20%20last_value%20%3D%20agent.value(cur_state)%0A%20%20%20%20%20%20%20%20buffer.compute_gae(last_value)%20%23%20Backwards%20pass%20to%20compute%20advantages%20and%20returns%0A%0A%20%20%20%20%20%20%20%20%23%20Returns%20the%20first%20state%20and%20the%20part-finished%20episode%20total%20used%20to%20initialise%20the%20next%20update%20cycle%0A%20%20%20%20%20%20%20%20return%20cur_state%2C%20total_reward%2C%20episode_rewards%0A%0A%20%20%20%20return%20(collect_rollout%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20This%20is%20the%20heart%20of%20PPO.%0A%0A%20%20%20%20The%20same%20rollout%20is%20reused%20for%20several%20epochs%2C%20so%20the%20data%20was%20gathered%20under%20an%20older%20policy%20%24%5Cpi_%7B%5Ctheta_%5Ctext%7Bold%7D%7D%24.%0A%20%20%20%20Importance%20sampling%20corrects%20for%20this%20with%20the%20ratio%20%24r_t(%5Ctheta)%20%3D%20%5Cpi_%5Ctheta(a_t%20%5Cvert%7B%7D%20s_t)%20%2F%20%5Cpi_%7B%5Ctheta_%5Ctext%7Bold%7D%7D(a_t%20%5Cvert%7B%7D%20s_t)%24%2C%20which%20we%20compute%20as%20a%20difference%20of%20log-probabilities%20inside%20an%20exponential.%0A%0A%20%20%20%20That%20correction%20is%20only%20trustworthy%20while%20the%20two%20policies%20remain%20close%2C%20so%20PPO%20takes%20the%20pessimistic%20(minimum)%20of%20the%20unclipped%20and%20clipped%20objectives%3A%0A%0A%20%20%20%20%24%24L%5E%7B%5Ctext%7BCLIP%7D%7D%20%3D%20%5Cmathbb%7BE%7D%5Cleft%5B%20%5Cmin%5Cleft(%20r_t%20A_t%2C%20%5C%3B%20%5Ctext%7Bclip%7D(r_t%2C%201-%5Cepsilon%2C%201%2B%5Cepsilon)%20A_t%20%5Cright)%20%5Cright%5D%24%24%0A%0A%20%20%20%20Once%20an%20update%20has%20pushed%20the%20ratio%20far%20enough%20in%20the%20direction%20that%20improves%20the%20objective%2C%20the%20clipped%20term%20becomes%20the%20smaller%20of%20the%20two%2C%20its%20gradient%20is%20zero%2C%20and%20there%20is%20no%20longer%20any%20incentive%20to%20move%20further.%0A%0A%20%20%20%20The%20total%20loss%20adds%20the%20critic's%20mean-squared%20error%20and%20subtracts%20an%20entropy%20bonus%20that%20discourages%20the%20policy%20from%20collapsing%20to%20a%20deterministic%20one%20too%20early.%0A%20%20%20%20We%20negate%20%24L%5E%7B%5Ctext%7BCLIP%7D%7D%24%20because%20the%20optimiser%20minimises%20rather%20than%20maximises.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(ActorCritic%2C%20RolloutBuffer%2C%20collect_rollout%2C%20gym%2C%20torch)%3A%0A%20%20%20%20def%20ppo_update_cycle(agent%20%3A%20ActorCritic%2C%20buffer%20%3A%20RolloutBuffer%2C%20env%20%3A%20gym.Env%2C%20optimiser%2C%20epochs%2C%20clip_eps%2C%20vf_coef%2C%20ent_coef%2C%20max_grad_norm%2C%20initial_state%2C%20initial_episode_reward)%3A%0A%20%20%20%20%20%20%20%20next_cycle_initial_state%2C%20next_cycle_episode_reward%2C%20episode_rewards%20%3D%20collect_rollout(agent%2C%20buffer%2C%20env%2C%20initial_state%2C%20initial_episode_reward)%0A%0A%20%20%20%20%20%20%20%20for%20epoch%20in%20range(epochs)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20buffer.new_epoch()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20minibatch%20in%20range(buffer.minibatches_per_epoch())%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Get%20data%20for%20loss%20function%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20states%2C%20actions%2C%20old_log_probs%2C%20advs%2C%20returns%2C%20old_values%20%3D%20buffer.get_minibatch(minibatch)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20new_log_probs%2C%20entropy%2C%20new_values%20%3D%20agent.evaluate(states%2C%20actions)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Compute%20the%20policy%20loss%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ratios%20%3D%20torch.exp(new_log_probs%20-%20old_log_probs)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20clipped_ratios%20%3D%20torch.clamp(ratios%2C%201%20-%20clip_eps%2C%201%20%2B%20clip_eps)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20L_1%20%3D%20ratios%20*%20advs%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20L_2%20%3D%20clipped_ratios%20*%20advs%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_loss%20%3D%20torch.min(L_1%2C%20L_2).mean()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Compute%20the%20value%20loss%20(MSE)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20value_loss%20%3D%200.5%20*%20(new_values%20-%20returns).pow(2).mean()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Compute%20the%20total%20loss%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss%20%3D%20-policy_loss%20%2B%20vf_coef%20*%20value_loss%20-%20ent_coef%20*%20entropy.mean()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20optimiser.zero_grad()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss.backward()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20torch.nn.utils.clip_grad_norm_(agent.parameters()%2C%20max_grad_norm)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20optimiser.step()%0A%0A%20%20%20%20%20%20%20%20return%20next_cycle_initial_state%2C%20next_cycle_episode_reward%2C%20episode_rewards%0A%0A%20%20%20%20return%20(ppo_update_cycle%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Training%20Loop%0A%20%20%20%20Tie%20everything%20together%20into%20one%20function%2C%20parameterised%20entirely%20by%20its%20arguments%20so%20the%20same%20code%20can%20be%20pointed%20at%20any%20continuous-control%20environment.%0A%20%20%20%20Each%20cycle%20collects%20a%20fresh%20rollout%20and%20then%20performs%20%60epochs_per_cycle%60%20passes%20of%20minibatch%20gradient%20descent%20over%20it.%0A%20%20%20%20Training%20stops%20early%20once%20the%20mean%20reward%20over%20the%20last%20100%20episodes%20exceeds%20%60reward_threshold%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(ActorCritic%2C%20RolloutBuffer%2C%20gym%2C%20np%2C%20ppo_update_cycle%2C%20torch)%3A%0A%20%20%20%20def%20train(%0A%20%20%20%20%20%20%20%20%20%20%20%20env_name%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20reward_threshold%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20neurons_per_layer%3D64%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20sigma_init%3D1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20buffer_size%20%3D%202048%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20optimiser_lr%20%3D%203e-4%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20minibatch_size%20%3D%2064%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20gamma%20%3D%200.9%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20lam%20%3D%200.95%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20epochs_per_cycle%20%3D%2010%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20clip_eps%20%3D%200.2%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vf_coef%20%3D%200.5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20ent_coef%20%3D%200.01%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20max_grad_norm%20%3D%200.5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20update_cycles%20%3D%201000%0A%20%20%20%20%20%20%20%20%20%20%20%20)%3A%0A%0A%20%20%20%20%20%20%20%20%23%20Setup%20the%20env%0A%20%20%20%20%20%20%20%20env%20%3D%20gym.make(env_name%2C%20render_mode%3D%22rgb_array%22)%0A%20%20%20%20%20%20%20%20state_dim%20%3D%20env.observation_space.shape%5B0%5D%0A%20%20%20%20%20%20%20%20action_dim%20%3D%20env.action_space.shape%5B0%5D%0A%20%20%20%20%20%20%20%20initial_state%2C%20_%20%3D%20env.reset()%0A%20%20%20%20%20%20%20%20initial_state%20%3D%20torch.tensor(initial_state%2C%20dtype%3Dtorch.float32)%0A%0A%20%20%20%20%20%20%20%20%23%20Build%20the%20agent%2C%20buffer%20and%20optimiser%0A%20%20%20%20%20%20%20%20agent%20%3D%20ActorCritic(state_dim%2C%20action_dim%2C%20neurons_per_layer%2C%20sigma_init)%0A%20%20%20%20%20%20%20%20buffer%20%3D%20RolloutBuffer(buffer_size%2C%20state_dim%2C%20action_dim%2C%20gamma%2C%20lam%2C%20minibatch_size)%0A%20%20%20%20%20%20%20%20optimiser%20%3D%20torch.optim.Adam(agent.parameters()%2C%20lr%3Doptimiser_lr)%0A%0A%20%20%20%20%20%20%20%20%23%20Every%20episode%20completed%20over%20the%20whole%20of%20training%2C%20in%20the%20order%20they%20finished%0A%20%20%20%20%20%20%20%20episode_rewards%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%23%20The%20running%20total%20of%20the%20episode%20in%20progress%2C%20threaded%20through%20the%20cycles%20alongside%20the%20state%0A%20%20%20%20%20%20%20%20unfinished_episode_reward%20%3D%200%0A%0A%20%20%20%20%20%20%20%20for%20cycle%20in%20range(update_cycles)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20initial_state%2C%20unfinished_episode_reward%2C%20cycle_episode_rewards%20%3D%20ppo_update_cycle(agent%2C%20buffer%2C%20env%2C%20optimiser%2C%20epochs_per_cycle%2C%20clip_eps%2C%20vf_coef%2C%20ent_coef%2C%20max_grad_norm%2C%20initial_state%2C%20unfinished_episode_reward)%0A%20%20%20%20%20%20%20%20%20%20%20%20episode_rewards.extend(cycle_episode_rewards)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20(cycle%2B1)%20%25%2010%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22Update%20cycle%20%7Bcycle%2B1%7D%20%7C%20Episode%20%7Blen(episode_rewards)%7D%20%7C%20Mean%20reward%20%7Bnp.mean(cycle_episode_rewards)%7D%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Compute%20the%20average%20reward%20over%20the%20last%20100%20episodes%20to%20verify%20reward%20threshold%20is%20met%0A%20%20%20%20%20%20%20%20%20%20%20%20start_idx%20%3D%20max(0%2C%20len(episode_rewards)%20-%20100)%0A%20%20%20%20%20%20%20%20%20%20%20%20mean_reward%20%3D%20np.mean(episode_rewards%5Bstart_idx%3A%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20mean_reward%20%3E%3D%20reward_threshold%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22Solved%20in%20%7Bcycle%2B1%7D%20update%20cycles%20with%20a%20mean%20reward%20of%20%7Bmean_reward%7D%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20break%0A%0A%20%20%20%20%20%20%20%20return%20env%2C%20agent%2C%20episode_rewards%0A%0A%20%20%20%20return%20(train%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%60Pendulum-v1%60%20makes%20a%20good%20smoke%20test.%0A%20%20%20%20It%20has%20a%20single%20action%20dimension%20and%20short%20episodes%2C%20so%20it%20trains%20quickly%20and%20exposes%20bugs%20before%20we%20attempt%20anything%20harder.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(train)%3A%0A%20%20%20%20env%2C%20agent%2C%20episode_rewards%20%3D%20train(%22Pendulum-v1%22%2C%20-150)%0A%20%20%20%20return%20agent%2C%20env%2C%20episode_rewards%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Plot%20the%20total%20reward%20of%20every%20completed%20episode%20against%20episode%20number%20to%20confirm%20that%20learning%20is%20actually%20progressing.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(episode_rewards%2C%20plt)%3A%0A%20%20%20%20%23%20Total%20reward%20against%20episode%20number%20over%20the%20course%20of%20training%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots()%0A%20%20%20%20_ax.scatter(range(len(episode_rewards))%2C%20episode_rewards%2C%20s%3D4)%0A%20%20%20%20_ax.set_xlabel(%22Episode%22)%0A%20%20%20%20_ax.set_ylabel(%22Total%20reward%22)%0A%20%20%20%20_ax.set_title(%22Training%20reward%20per%20episode%22)%0A%20%20%20%20plt.close(_fig)%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20During%20training%20we%20sample%20from%20%24%5Cmathcal%7BN%7D(%5Cmu(s)%2C%20%5Csigma%5E2)%24%20in%20order%20to%20explore.%0A%20%20%20%20To%20evaluate%20the%20learned%20policy%20we%20stop%20exploring%20and%20take%20the%20mean%20%24%5Cmu(s)%24%20directly%2C%20which%20is%20the%20continuous%20analogue%20of%20the%20greedy%20action.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(ActorCritic%2C%20gym%2C%20torch)%3A%0A%20%20%20%20%23%20Pure%20exploitation%3A%20always%20take%20the%20mean%20of%20the%20policy%20rather%20than%20sampling%20from%20it%0A%20%20%20%20def%20evaluate(agent%20%3A%20ActorCritic%2C%20env%20%3A%20gym.Env%2C%20episodes)%3A%0A%20%20%20%20%20%20%20%20episode_rewards%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20for%20_%20in%20range(episodes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20_%20%3D%20env.reset()%0A%20%20%20%20%20%20%20%20%20%20%20%20done%20%3D%20False%0A%20%20%20%20%20%20%20%20%20%20%20%20total_reward%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20while%20not%20done%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20action%20%3D%20agent.greedy_action(torch.tensor(state%2C%20dtype%3Dtorch.float32))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20reward%2C%20terminated%2C%20truncated%2C%20_%20%3D%20env.step(action.numpy())%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20done%20%3D%20terminated%20or%20truncated%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20total_reward%20%2B%3D%20reward%0A%20%20%20%20%20%20%20%20%20%20%20%20episode_rewards.append(total_reward)%0A%0A%20%20%20%20%20%20%20%20return%20episode_rewards%0A%0A%20%20%20%20return%20(evaluate%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20A%20single%20episode%20is%20noisy%20because%20the%20initial%20state%20is%20random%2C%20so%20we%20average%20over%20many%20episodes%20and%20report%20the%20standard%20deviation%20alongside%20the%20mean.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(agent%2C%20env%2C%20evaluate%2C%20np)%3A%0A%20%20%20%20%23%20Average%20over%20several%20greedy%20episodes%20to%20get%20a%20true%20measure%20of%20the%20learned%20policy%0A%20%20%20%20greedy_rewards%20%3D%20evaluate(agent%2C%20env%2C%20episodes%3D50)%0A%20%20%20%20print(f%22Greedy%20average%20reward%20over%20%7Blen(greedy_rewards)%7D%20episodes%3A%20%7Bnp.mean(greedy_rewards)%3A.2f%7D%20%2B%2F-%20%7Bnp.std(greedy_rewards)%3A.2f%7D%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Finally%2C%20record%20the%20frames%20of%20one%20greedy%20episode%20and%20stitch%20them%20together%20into%20an%20animation%20so%20that%20the%20learned%20policy%20can%20be%20watched%20directly.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(ActorCritic%2C%20Image%2C%20base64%2C%20gym%2C%20io%2C%20mo%2C%20torch)%3A%0A%20%20%20%20def%20animate_episode(agent%20%3A%20ActorCritic%2C%20env%20%3A%20gym.Env%2C%20interval%3D40)%3A%0A%20%20%20%20%20%20%20%20frames%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20state%2C%20_%20%3D%20env.reset()%0A%20%20%20%20%20%20%20%20done%20%3D%20False%0A%20%20%20%20%20%20%20%20total_reward%20%3D%200%0A%0A%20%20%20%20%20%20%20%20while%20not%20done%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Render%20before%20stepping%20so%20the%20very%20first%20state%20is%20captured%0A%20%20%20%20%20%20%20%20%20%20%20%20frames.append(env.render())%0A%20%20%20%20%20%20%20%20%20%20%20%20action%20%3D%20agent.greedy_action(torch.tensor(state%2C%20dtype%3Dtorch.float32))%0A%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20reward%2C%20terminated%2C%20truncated%2C%20_%20%3D%20env.step(action.numpy())%0A%20%20%20%20%20%20%20%20%20%20%20%20done%20%3D%20terminated%20or%20truncated%0A%20%20%20%20%20%20%20%20%20%20%20%20total_reward%20%2B%3D%20reward%0A%0A%20%20%20%20%20%20%20%20%23%20Encode%20the%20frames%20as%20a%20looping%20GIF.%20Unlike%20matplotlib's%20%60to_jshtml%60%2C%20this%20needs%20no%0A%20%20%20%20%20%20%20%20%23%20JavaScript%2C%20so%20it%20also%20displays%20in%20a%20static%20HTML%20export%20of%20the%20notebook.%0A%20%20%20%20%20%20%20%20images%20%3D%20%5BImage.fromarray(frame)%20for%20frame%20in%20frames%5D%0A%20%20%20%20%20%20%20%20gif%20%3D%20io.BytesIO()%0A%20%20%20%20%20%20%20%20images%5B0%5D.save(gif%2C%20format%3D%22GIF%22%2C%20save_all%3DTrue%2C%20append_images%3Dimages%5B1%3A%5D%2C%20duration%3Dinterval%2C%20loop%3D0)%0A%20%20%20%20%20%20%20%20gif_b64%20%3D%20base64.b64encode(gif.getvalue()).decode()%0A%20%20%20%20%20%20%20%20return%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.md(f%22Greedy%20episode%2C%20total%20reward%3A%20%7Btotal_reward%3A.2f%7D%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.Html(f'%3Cimg%20src%3D%22data%3Aimage%2Fgif%3Bbase64%2C%7Bgif_b64%7D%22%20alt%3D%22Greedy%20episode%20of%20the%20trained%20agent%22%3E')%2C%0A%20%20%20%20%20%20%20%20%5D)%0A%0A%20%20%20%20return%20(animate_episode%2C)%0A%0A%0A%40app.cell%0Adef%20_(agent%2C%20animate_episode%2C%20env)%3A%0A%20%20%20%20animate_episode(agent%2C%20env)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20pendulum%20is%20hence%20swung%20up%20and%20balanced%20using%20PPO.%0A%20%20%20%20Since%20%60train()%60%20is%20configured%20entirely%20through%20its%20arguments%2C%20the%20same%20notebook%20can%20be%20pointed%20at%20a%20harder%20environment%20such%20as%20%60BipedalWalker-v3%60%20by%20changing%20%60env_name%60%20and%20%60reward_threshold%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A