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%20random%0A%20%20%20%20from%20collections%20import%20deque%0A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20gymnasium%20as%20gym%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%20%20%20%20import%20torch.nn.functional%20as%20F%0A%0A%20%20%20%20return%20F%2C%20deque%2C%20gym%2C%20mo%2C%20nn%2C%20np%2C%20plt%2C%20random%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%20Basic%20Solution%20to%20The%20Cart-Pole%20Problem%20using%20a%20Deep%20Q-Network%0A%20%20%20%20!%5B%5D(https%3A%2F%2Fgymnasium.farama.org%2F_images%2Fcart_pole.gif)%0A%0A%20%20%20%20Basic%20code%20to%20implement%20a%20Deep%20Q-Network%20to%20solve%20the%20Cart-Pole%20problem%20is%20given.%0A%0A%20%20%20%20A%20pole%20is%20balancing%20on%20a%20cart%20which%20can%20move%20left%20and%20right.%0A%20%20%20%20The%20only%20thing%20the%20cart%20can%20do%20to%20balance%20the%20pole%20is%20apply%20a%20fixed%20force%20to%20the%20left%20or%20right.%0A%20%20%20%20If%20the%20pole%20hits%20the%20ground%20the%20cart's%20family%20will%20be%20killed.%0A%20%20%20%20We%20aid%20the%20cart%20by%20setting%20up%20a%20simulation%20to%20help%20him%20practice%20using%20the%20highly%20limited%203%20digit%20number%20of%20neurons%20his%20brain%20has.%0A%0A%20%20%20%20Luckily%20for%20the%20cart%2C%20this%20is%20a%20well%20established%20problem%20in%20reinforcement%20learning%2C%20so%20we%20can%20simply%20import%20a%20pre-existing%20physics%20simulation%20of%20his%20current%20unfortunate%20predicament.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(gym%2C%20np%2C%20random%2C%20torch)%3A%0A%20%20%20%20%23%20For%20determinism%0A%20%20%20%20global_seed%20%3D%2067%0A%0A%20%20%20%20%23%20Seed%20every%20random%20number%20generator%20up%20front%20so%20that%20everything%20below%20(network%20initialisation%2C%0A%20%20%20%20%23%20replay%20sampling%2C%20epsilon-greedy%20exploration)%20is%20reproducible%0A%20%20%20%20torch.manual_seed(global_seed)%0A%20%20%20%20np.random.seed(global_seed)%0A%20%20%20%20random.seed(global_seed)%0A%0A%20%20%20%20%23%20Initialize%20the%20CartPole%20environment%0A%20%20%20%20env%20%3D%20gym.make(%22CartPole-v1%22)%0A%0A%20%20%20%20%23%20Get%20the%20dimensions%20of%20the%20state%20space%20and%20the%20action%20space%0A%20%20%20%20state_dim%20%3D%20env.observation_space.shape%5B0%5D%0A%20%20%20%20action_dim%20%3D%20env.action_space.n%0A%0A%20%20%20%20print(f%22State%20space%20dimension%3A%20%7Bstate_dim%7D%22)%0A%20%20%20%20print(f%22Action%20space%20dimension%3A%20%7Baction_dim%7D%22)%0A%20%20%20%20return%20action_dim%2C%20env%2C%20global_seed%2C%20state_dim%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%20Now%20we%20create%20a%20model%20of%20the%20cart's%20brain%20by%20inheriting%20%60torch.nn.Module%60.%0A%20%20%20%20He%20has%204%20sensory%20inputs%20(state%20space)%20relevant%20to%20this%20training%20simulation%3A%0A%0A%20%20%20%20*%20His%20position%0A%20%20%20%20*%20His%20speed%0A%20%20%20%20*%20The%20pole's%20angle%20relative%20to%20him%0A%20%20%20%20*%20The%20pole's%20angular%20velocity%0A%0A%20%20%20%20He%20then%20has%202%20layers%20of%2064%20neurons%20each%20(shame)%2C%20and%20then%20a%20final%20output%20layer%20to%20predict%20the%20reward%20if%20he%20applies%20a%20force%20left%20or%20applies%20a%20force%20right.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20nn)%3A%0A%20%20%20%20class%20QNetwork(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20state_dim%2C%20action_dim)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super(QNetwork%2C%20self).__init__()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Map%20from%20layer%201%20to%202%3A%20Takes%20the%204%20state%20variables%2C%20outputs%20to%2064%20hidden%20neurons%0A%20%20%20%20%20%20%20%20%20%20%20%20self.fc1%20%3D%20nn.Linear(state_dim%2C%2064)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Map%20from%20layer%202%20to%203%3A%20Hidden%20layer%20to%20hidden%20layer%0A%20%20%20%20%20%20%20%20%20%20%20%20self.fc2%20%3D%20nn.Linear(64%2C%2064)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Map%20from%20layer%203%20to%204%3A%20Takes%20the%2064%20hidden%20neurons%2C%20outputs%20the%202%20Q-values%0A%20%20%20%20%20%20%20%20%20%20%20%20self.fc3%20%3D%20nn.Linear(64%2C%20action_dim)%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20state)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20We%20use%20ReLU(x)%3Dmax(x%2C0)%20activation%20for%20the%20hidden%20layers%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20F.relu(self.fc1(state))%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20F.relu(self.fc2(x))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20No%20activation%20function%20on%20the%20final%20layer%20because%20Q-values%20can%20theoretically%20be%20any%20real%20number%20(positive%20or%20negative)%0A%20%20%20%20%20%20%20%20%20%20%20%20q_values%20%3D%20self.fc3(x)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20q_values%0A%0A%20%20%20%20return%20(QNetwork%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%20Give%20him%20a%20memory%20of%2010000%20past%20experiences.%0A%20%20%20%20Once%20his%20memory%20is%20full%20he%20forgets%20the%20oldest%20experience.%0A%20%20%20%20To%20aid%20in%20his%20learning%20we%20allow%20him%20to%20recall%20up%20to%2064%20past%20experiences%20randomly%20to%20learn%20from%20in%20his%20training.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(deque%2C%20np%2C%20random)%3A%0A%20%20%20%20%23%20Define%20the%20class%0A%20%20%20%20class%20ReplayBuffer%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20capacity%3D10000)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20deque%20automatically%20removes%20oldest%20items%20if%20capacity%20is%20exceeded%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer%20%3D%20deque(maxlen%3Dcapacity)%0A%0A%20%20%20%20%20%20%20%20def%20push(self%2C%20state%2C%20action%2C%20reward%2C%20next_state%2C%20done)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Store%20the%205-tuple%20experience%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer.append((state%2C%20action%2C%20reward%2C%20next_state%2C%20done))%0A%0A%20%20%20%20%20%20%20%20def%20sample(self%2C%20batch_size)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Randomly%20select%20a%20batch%20of%20experiences%0A%20%20%20%20%20%20%20%20%20%20%20%20batch%20%3D%20random.sample(self.buffer%2C%20batch_size)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Unzip%20the%20batch%20into%20separate%20lists%2Farrays%20for%20PyTorch%0A%20%20%20%20%20%20%20%20%20%20%20%20states%2C%20actions%2C%20rewards%2C%20next_states%2C%20dones%20%3D%20map(np.array%2C%20zip(*batch))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20states%2C%20actions%2C%20rewards%2C%20next_states%2C%20dones%0A%0A%20%20%20%20%20%20%20%20def%20__len__(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20len(self.buffer)%0A%0A%20%20%20%20%23%20Instantiate%20it%0A%20%20%20%20memory%20%3D%20ReplayBuffer(capacity%3D10000)%0A%20%20%20%20return%20(memory%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%20In%20his%20practice%2C%20it%20is%20crucial%20that%20he%20explore%20new%20strategies%20every%20now%20and%20then.%20To%20do%20this%20we%20use%20%24%5Cepsilon%24-greedy%20action%20selection.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(random%2C%20torch)%3A%0A%20%20%20%20def%20select_action(state%2C%20q_net%2C%20epsilon%2C%20action_dim)%3A%0A%20%20%20%20%20%20%20%20if%20random.random()%20%3C%20epsilon%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Explore%3A%20choose%20a%20random%20action%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20random.randrange(action_dim)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Exploit%3A%20choose%20the%20action%20with%20the%20highest%20Q-value%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Convert%20the%20numpy%20state%20vector%20into%20a%20pytorch%20tensor%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Unsqueeze%20changes%20the%20tensor%20from%201D%20to%202D%0A%20%20%20%20%20%20%20%20%20%20%20%20state_tensor%20%3D%20torch.FloatTensor(state).unsqueeze(0)%20%20%23%20Shape%3A%20(1%2C%204)%0A%20%20%20%20%20%20%20%20%20%20%20%20with%20torch.no_grad()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_values%20%3D%20q_net(state_tensor)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20q_values.argmax().item()%0A%0A%20%20%20%20return%20(select_action%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%20Make%20a%20copy%20of%20his%20brain.%20His%20main%20brain%20is%20the%20policy%20network%2C%20and%20his%20secondary%20brain%20is%20the%20target%20network.%0A%20%20%20%20His%20second%20brain%20is%20kept%20frozen%20and%20only%20updated%20(by%20copying%20his%20main%20brain)%20every%20%24n%24%20iterations.%0A%20%20%20%20The%20policy%20network%20guesses%20the%20action%20value%2C%20the%20target%20network%20(frozen%20copy)%20evaluates%20the%20actual%20return%20of%20that%20action%2C%20as%20well%20as%20the%20discounted%20return%20of%20the%20guess%20of%20the%20next%20action%20value.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNetwork%2C%20action_dim%2C%20state_dim)%3A%0A%20%20%20%20%23%20Initialize%20both%20networks%0A%20%20%20%20policy_net%20%3D%20QNetwork(state_dim%2C%20action_dim)%0A%20%20%20%20target_net%20%3D%20QNetwork(state_dim%2C%20action_dim)%0A%0A%20%20%20%20%23%20Copy%20the%20completely%20random%20initial%20weights%20from%20policy%20to%20target%0A%20%20%20%20target_net.load_state_dict(policy_net.state_dict())%0A%0A%20%20%20%20%23%20Set%20target%20network%20to%20evaluation%20mode%20(turns%20off%20gradients%2Ftraining%20behaviors)%0A%20%20%20%20target_net.eval()%0A%20%20%20%20return%20policy_net%2C%20target_net%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%20Setup%20the%20loss%20function.%20This%20is%20a%20measure%20of%20how%20'surprised'%20the%20policy%20network%20is%20by%20its%20action%20value%20guess%20versus%20the%20target%20networks%20more%20accurate%20guess%20one%20step%20in%20the%20future.%0A%20%20%20%20We%20aim%20for%20the%20cart%20to%20never%20be%20surprised%20by%20the%20outcome%20of%20an%20action.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20torch)%3A%0A%20%20%20%20def%20compute_loss(batch%2C%20policy_net%2C%20target_net%2C%20gamma%3D0.99)%3A%0A%20%20%20%20%20%20%20%20%23%20Unpack%20the%20batch%20of%20experiences%0A%20%20%20%20%20%20%20%20states%2C%20actions%2C%20rewards%2C%20next_states%2C%20dones%20%3D%20batch%0A%0A%20%20%20%20%20%20%20%20%23%20Convert%20numpy%20arrays%20to%20PyTorch%20tensors%0A%20%20%20%20%20%20%20%20states%20%3D%20torch.FloatTensor(states)%0A%20%20%20%20%20%20%20%20actions%20%3D%20torch.LongTensor(actions).unsqueeze(1)%20%23%20Reshape%20for%20gather%0A%20%20%20%20%20%20%20%20rewards%20%3D%20torch.FloatTensor(rewards)%0A%20%20%20%20%20%20%20%20next_states%20%3D%20torch.FloatTensor(next_states)%0A%20%20%20%20%20%20%20%20dones%20%3D%20torch.FloatTensor(dones)%0A%0A%20%20%20%20%20%20%20%20%23%201.%20PREDICTIONS%3A%20Get%20the%20Q-values%20from%20the%20Policy%20Network%20for%20the%20actions%20actually%20taken%0A%20%20%20%20%20%20%20%20%23%20.gather()%20acts%20like%20a%20router%2C%20plucking%20out%20the%20exact%20Q-value%20for%20the%20chosen%20action%0A%20%20%20%20%20%20%20%20current_q_values%20%3D%20policy_net(states).gather(1%2C%20actions).squeeze(1)%0A%0A%20%20%20%20%20%20%20%20%23%202.%20TARGETS%3A%20Get%20the%20maximum%20Q-values%20for%20the%20next%20states%20from%20the%20Target%20Network%0A%20%20%20%20%20%20%20%20with%20torch.no_grad()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20.max(1)%5B0%5D%20gets%20the%20highest%20Q-value%20across%20the%20action%20dimension%0A%20%20%20%20%20%20%20%20%20%20%20%20next_q_values%20%3D%20target_net(next_states).max(1)%5B0%5D%0A%0A%20%20%20%20%20%20%20%20%23%20Calculate%20the%20Bellman%20targets%0A%20%20%20%20%20%20%20%20target_q_values%20%3D%20rewards%20%2B%20(gamma%20*%20next_q_values%20*%20(1%20-%20dones))%0A%0A%20%20%20%20%20%20%20%20%23%203.%20LOSS%3A%20Compare%20them%20using%20Mean%20Squared%20Error%0A%20%20%20%20%20%20%20%20loss%20%3D%20F.mse_loss(current_q_values%2C%20target_q_values)%0A%0A%20%20%20%20%20%20%20%20return%20loss%0A%0A%20%20%20%20return%20(compute_loss%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%20Use%20the%20Adam%20algorithm%20to%20adjust%20his%20neuron%20weights%20and%20biases%20to%20minimise%20the%20loss%20function.%0A%20%20%20%20In%20simpler%20terms%2C%20we%20alter%20his%20brain%20to%20minimise%20his%20surprise%20of%20the%20outcome%20of%20an%20action.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(policy_net%2C%20torch)%3A%0A%20%20%20%20%23%20Set%20up%20the%20optimiser%0A%20%20%20%20optimizer%20%3D%20torch.optim.Adam(policy_net.parameters()%2C%20lr%3D1e-3)%0A%0A%20%20%20%20%23%20Hyperparameters%0A%20%20%20%20num_episodes%20%3D%20500%0A%20%20%20%20batch_size%20%3D%2064%0A%20%20%20%20gamma%20%3D%200.99%0A%20%20%20%20epsilon%20%3D%201.0%20%23%20Initial%20value%20of%20epsilon%3B%20it%20is%20decayed%20during%20training%0A%20%20%20%20epsilon_min%20%3D%200.01%0A%20%20%20%20epsilon_decay%20%3D%200.990%0A%20%20%20%20reward_threshold%20%3D%20475%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20batch_size%2C%0A%20%20%20%20%20%20%20%20epsilon%2C%0A%20%20%20%20%20%20%20%20epsilon_decay%2C%0A%20%20%20%20%20%20%20%20epsilon_min%2C%0A%20%20%20%20%20%20%20%20gamma%2C%0A%20%20%20%20%20%20%20%20num_episodes%2C%0A%20%20%20%20%20%20%20%20optimizer%2C%0A%20%20%20%20%20%20%20%20reward_threshold%2C%0A%20%20%20%20)%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%20Before%20setting%20up%20the%20main%20training%20loop%2C%20we%20create%20a%20function%20to%20evaluate%20the%20cart's%20training%20progress.%0A%20%20%20%20We%20set%20%24%5Cepsilon%20%3D%200%24%20to%20stop%20exploration%2C%20and%20make%20the%20cart%20only%20exploit%20its%20knowledge.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(action_dim%2C%20env%2C%20global_seed%2C%20policy_net%2C%20select_action)%3A%0A%20%20%20%20def%20run_greedy_episode(use_seed%20%3D%20False)%3A%0A%20%20%20%20%20%20%20%20if%20use_seed%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20_%20%3D%20env.reset(seed%3Dglobal_seed)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20_%20%3D%20env.reset()%0A%0A%20%20%20%20%20%20%20%20total_reward%20%3D%200%0A%20%20%20%20%20%20%20%20done%20%3D%20False%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%20action%20%3D%20select_action(state%2C%20policy_net%2C%200%2C%20action_dim)%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)%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%20%23%20Not%20necessary%20to%20sync%20the%20target%20network%2C%20update%20memory%20or%20train%0A%20%20%20%20%20%20%20%20%20%20%20%20state%20%3D%20next_state%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%20return%20total_reward%0A%0A%20%20%20%20return%20(run_greedy_episode%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%20Main%20training%20loop.%0A%20%20%20%20We%20fully%20simulate%20each%20episode%2C%20adding%20each%20experience%20to%20memory.%0A%20%20%20%20Once%20enough%20memories%20have%20accrued%2C%20we%20randomly%20sample%20a%20%60batch_size%60%20number%20of%20memories%20and%20optimise%20the%20network%20parameters%20based%20on%20those.%0A%0A%20%20%20%20If%20the%20reward%20of%20one%20of%20these%20episodes%20is%20close%20to%20perfect%2C%20we%20%60run_greedy_episode()%60%20a%20number%20of%20times%20to%20verify%20those%20network%20parameters%20are%20indeed%20a%20solution.%0A%20%20%20%20If%20they%20are%2C%20we%20stop%20training.%0A%0A%20%20%20%20The%20loop%20lives%20in%20a%20function%20so%20that%20its%20working%20variables%20(and%20the%20decayed%20value%20of%20%24%5Cepsilon%24)%20stay%20local%20to%20it%2C%20which%20marimo%20requires%20because%20every%20global%20name%20may%20only%20be%20defined%20in%20one%20cell.%0A%20%20%20%20It%20also%20returns%20the%20reward%20and%20%24%5Cepsilon%24%20of%20every%20episode%20so%20that%20we%20can%20plot%20them%20below.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20action_dim%2C%0A%20%20%20%20batch_size%2C%0A%20%20%20%20compute_loss%2C%0A%20%20%20%20env%2C%0A%20%20%20%20epsilon%2C%0A%20%20%20%20epsilon_decay%2C%0A%20%20%20%20epsilon_min%2C%0A%20%20%20%20gamma%2C%0A%20%20%20%20global_seed%2C%0A%20%20%20%20memory%2C%0A%20%20%20%20np%2C%0A%20%20%20%20num_episodes%2C%0A%20%20%20%20optimizer%2C%0A%20%20%20%20policy_net%2C%0A%20%20%20%20reward_threshold%2C%0A%20%20%20%20run_greedy_episode%2C%0A%20%20%20%20select_action%2C%0A%20%20%20%20target_net%2C%0A)%3A%0A%20%20%20%20def%20train()%3A%0A%20%20%20%20%20%20%20%20total_training_steps%20%3D%200%0A%20%20%20%20%20%20%20%20%23%20The%20current%20exploration%20rate%2C%20decayed%20after%20every%20episode%0A%20%20%20%20%20%20%20%20eps%20%3D%20epsilon%0A%20%20%20%20%20%20%20%20episode_rewards%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20episode_epsilons%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20for%20episode%20in%20range(num_episodes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20episode%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state%2C%20_%20%3D%20env.reset(seed%3Dglobal_seed)%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%20state%2C%20_%20%3D%20env.reset()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Reset%20to%20random%20initial%20conditions%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%20done%20%3D%20False%0A%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%20%23%20Choose%20the%20action%20with%20the%20highest%20predicted%20immediate%20reward%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20action%20%3D%20select_action(state%2C%20policy_net%2C%20eps%2C%20action_dim)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Step%20the%20physics%20simulation%20forward%20in%20time%20to%20get%20the%20actual%20reward%0A%20%20%20%20%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)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20If%20the%20network%20fails%2C%20or%20makes%20it%20to%20the%20end%20of%20the%20500%20iterations%2C%20terminate%20the%20while%20loop%0A%20%20%20%20%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%20%20%20%20%20%23%20Add%20the%20current%20experience%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20memory.push(state%2C%20action%2C%20reward%2C%20next_state%2C%20terminated)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state%20%3D%20next_state%0A%20%20%20%20%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%20%20%20%20%23%20Once%20enough%20memory%20has%20been%20accrued%20we%20sample%20the%20memory%20to%20make%20the%20model%20learn%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20len(memory)%20%3E%3D%205*batch_size%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch%20%3D%20memory.sample(batch_size)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss%20%3D%20compute_loss(batch%2C%20policy_net%2C%20target_net%2C%20gamma)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Improve%20the%20policy%20network's%20parameters%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20optimizer.zero_grad()%0A%20%20%20%20%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%20%20%20%20%20optimizer.step()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20total_training_steps%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20total_training_steps%20%25%20500%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20target_net.load_state_dict(policy_net.state_dict())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Decay%20epsilon%20-%20reduce%20the%20amount%20of%20exploration%20with%20time%20to%20some%20minimum%20value%0A%20%20%20%20%20%20%20%20%20%20%20%20eps%20%3D%20max(epsilon_min%2C%20epsilon_decay%20*%20eps)%0A%0A%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%20episode_epsilons.append(eps)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Log%20the%20current%20approximate%20training%20status%20every%20few%20episodes%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20(episode%20%2B%201)%20%25%2020%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22Episode%20%7Bepisode%20%2B%201%7D%20%7C%20Total%20Reward%3A%20%7Btotal_reward%3A.1f%7D%20%7C%20Epsilon%3A%20%7Beps%3A.2f%7D%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20If%20the%20reward%20is%20greater%20than%20some%20threshold%2C%20pause%20training%20to%20verify%20if%20the%20solution%20is%20good%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20total_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%22Found%20potential%20solution%20at%20episode%20%7Bepisode%20%2B%201%7D%2C%20total%20reward%20of%20%7Btotal_reward%7D%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Set%20epsilon%20to%20zero%20and%20run%20n%20times%20for%20testing%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20scores%20%3D%20%5Brun_greedy_episode()%20for%20_%20in%20range(10)%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_score%20%3D%20np.mean(scores)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Stop%20training%20if%20solution%20is%20correct%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20avg_score%20%3E%3D%20reward_threshold%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22Confirmed%20solution%20at%20episode%20%7Bepisode%20%2B%201%7D%2C%20averaging%20a%20reward%20of%20%7Bavg_score%7D%22)%0A%20%20%20%20%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%20episode_rewards%2C%20episode_epsilons%0A%0A%20%20%20%20episode_rewards%2C%20episode_epsilons%20%3D%20train()%0A%20%20%20%20return%20episode_epsilons%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%20training%20episode%20(left)%2C%20together%20with%20the%20exploration%20rate%20%24%5Cepsilon%24%20(right)%2C%20to%20see%20the%20cart%20move%20from%20exploring%20to%20exploiting.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(episode_epsilons%2C%20episode_rewards%2C%20np%2C%20plt%2C%20reward_threshold)%3A%0A%20%20%20%20_window%20%3D%2020%0A%20%20%20%20_episodes%20%3D%20np.arange(1%2C%20len(episode_rewards)%20%2B%201)%0A%0A%20%20%20%20_fig%2C%20(_ax_reward%2C%20_ax_eps)%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(11%2C%203.8))%0A%0A%20%20%20%20_ax_reward.scatter(_episodes%2C%20episode_rewards%2C%20s%3D10%2C%20color%3D%22dodgerblue%22%2C%20alpha%3D0.6%2C%20label%3D%22Episode%20reward%22)%0A%20%20%20%20if%20len(episode_rewards)%20%3E%3D%20_window%3A%0A%20%20%20%20%20%20%20%20_avg%20%3D%20np.convolve(episode_rewards%2C%20np.ones(_window)%20%2F%20_window%2C%20mode%3D%22valid%22)%0A%20%20%20%20%20%20%20%20_ax_reward.plot(_episodes%5B_window%20-%201%3A%5D%2C%20_avg%2C%20color%3D%22red%22%2C%20linewidth%3D2%2C%20label%3Df%22Average%20of%20last%20%7B_window%7D%20episodes%22)%0A%20%20%20%20_ax_reward.axhline(reward_threshold%2C%20color%3D%22grey%22%2C%20linestyle%3D%22%3A%22%2C%20label%3Df%22Threshold%20(%7Breward_threshold%7D)%22)%0A%20%20%20%20_ax_reward.set_xlabel(%22Episode%20%23%22)%0A%20%20%20%20_ax_reward.set_ylabel(%22Total%20reward%22)%0A%20%20%20%20_ax_reward.set_title(%22DQN%20training%20progress%20on%20CartPole-v1%22)%0A%20%20%20%20_ax_reward.grid(True%2C%20linestyle%3D%22--%22%2C%20alpha%3D0.5)%0A%20%20%20%20_ax_reward.legend(loc%3D%22upper%20left%22)%0A%0A%20%20%20%20_ax_eps.plot(_episodes%2C%20episode_epsilons%2C%20color%3D%22darkorange%22%2C%20linewidth%3D2)%0A%20%20%20%20_ax_eps.set_xlabel(%22Episode%20%23%22)%0A%20%20%20%20_ax_eps.set_ylabel(r%22%24%5Cepsilon%24%22)%0A%20%20%20%20_ax_eps.set_title(r%22Exploration%20rate%20%24%5Cepsilon%24%22)%0A%20%20%20%20_ax_eps.grid(True%2C%20linestyle%3D%22--%22%2C%20alpha%3D0.5)%0A%0A%20%20%20%20_fig.tight_layout()%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%20Now%20we%20verify%20if%20the%20cart%20has%20learned%20enough%20to%20save%20his%20family's%20life.%0A%20%20%20%20We%20stop%20exploration%20by%20greedily%20choosing%20the%20best%20action%20according%20to%20the%20network's%20approximation%20of%20%24Q(s%2Ca)%24%2C%20with%20%24%5Cepsilon%20%3D%200%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(env%2C%20episode_rewards%2C%20np%2C%20run_greedy_episode)%3A%0A%20%20%20%20%23%20Referencing%20%60episode_rewards%60%20makes%20this%20cell%20run%20after%20the%20training%20loop%0A%20%20%20%20print(f%22Evaluating%20the%20policy%20obtained%20after%20%7Blen(episode_rewards)%7D%20training%20episodes%22)%0A%0A%20%20%20%20n_evals%20%3D%20100%0A%20%20%20%20eval_rewards%20%3D%20np.zeros(n_evals)%0A%0A%20%20%20%20for%20eval_episode%20in%20range(n_evals)%3A%0A%20%20%20%20%20%20%20%20eval_rewards%5Beval_episode%5D%20%3D%20run_greedy_episode(eval_episode%20%3D%3D%200)%0A%0A%20%20%20%20avg_reward%20%3D%20np.mean(eval_rewards)%0A%20%20%20%20std_reward%20%3D%20np.std(eval_rewards)%0A%0A%20%20%20%20print(f%22Average%20reward%20over%20%7Bn_evals%7D%20evaluations%3A%20%7Bavg_reward%3A.2f%7D%22)%0A%20%20%20%20print(f%22Standard%20deviation%20of%20rewards%3A%20%7Bstd_reward%3A.2f%7D%22)%0A%0A%20%20%20%20env.close()%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%20%20%20%20The%20cart%20is%20ready%20to%20be%20tested%20to%20save%20his%20family's%20life.%0A%20%20%20%20%20%20%20%20His%20training%20was%20successful.%0A%20%20%20%20%20%20%20%20%22%22%22)%2C%0A%20%20%20%20%20%20%20%20mo.image(src%3Dstr(mo.notebook_dir()%20%2F%20%22assets%22%20%2F%20%22Cart-Family.jpg%22)%2C%20alt%3D%22Cart-Family.jpg%22%2C%20width%3D640)%2C%0A%20%20%20%20%5D)%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
dddfa95f81e3fdf4f8cf93160c23cfab