import%20marimo%0A%0A__generated_with%20%3D%20%220.19.2%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22%2C%20app_title%3D%22RL4%3A%20Dynamic%20Programming%22)%0A%0Awith%20app.setup(hide_code%3DTrue)%3A%0A%20%20%20%20%23%20Initialization%20code%20that%20runs%20before%20all%20other%20cells%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20import%20gymnasium%20as%20gym%0A%20%20%20%20from%20gymnasium.envs.toy_text.frozen_lake%20import%20generate_random_map%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20seaborn%20as%20sns%0A%20%20%20%20import%20warnings%20%0A%20%20%20%20import%20functools%0A%20%20%20%20warnings.filterwarnings(%22ignore%22%2C%20category%3DRuntimeWarning)%0A%20%20%20%20warnings.filterwarnings(%22ignore%22%2C%20category%3DUserWarning)%0A%20%20%20%20np.random.seed(308)%0A%0A%20%20%20%20def%20running_avg(data%2C%20window_size)%3A%0A%20%20%20%20%20%20%20%20res%20%3D%20np.zeros(len(data)-window_size)%0A%20%20%20%20%20%20%20%20sum%20%3D%200%0A%20%20%20%20%20%20%20%20for%20i%20in%20range(window_size)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20sum%20%2B%3D%20data%5Bi%5D%0A%20%20%20%20%20%20%20%20for%20i%20in%20range(len(data)-window_size)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20res%5Bi%5D%20%3D%20sum%20%2F%20window_size%0A%20%20%20%20%20%20%20%20%20%20%20%20sum%20-%3D%20data%5Bi%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20sum%20%2B%3D%20data%5Bi%20%2B%20window_size%5D%0A%20%20%20%20%20%20%20%20return%20res%0A%0A%20%20%20%20def%20configure_seaborn(**kwargs)%3A%0A%20%20%20%20%20%20%20%20sns.set_context(%22notebook%22)%0A%20%20%20%20%20%20%20%20sns.set_theme(sns.plotting_context(%22notebook%22%2C%20font_scale%3D1)%2C%20style%3D%22whitegrid%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20rc%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20grid%20activation%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22axes.grid%22%3A%20False%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20grid%20appearance%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22grid.color%22%3A%20%22%23BFBFBF%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22axes.edgecolor%22%3A%20%22%23BFBFBF%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22axes.linewidth%22%3A%200.8%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22grid.linewidth%22%3A%200.5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22grid.alpha%22%3A%200.4%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20'font.family'%3A'sans-serif'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C)%0A%0A%20%20%20%20%20%20%20%20palette%20%3D%20%5B%22%233D405B%22%2C%20%22%23E08042%22%2C%20%22%2354AB69%22%2C%20%22%23CE2F49%22%2C%20%22%23A26EBF%22%2C%20%22%2375523A%22%2C%20%22%23D12AA2%22%2C%20%22%23E0D06F%22%2C%20%22%236F9AA7%22%2C%20%22%233359C4%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22%2376455B%22%5D%0A%0A%20%20%20%20%20%20%20%20sns.set_palette(palette)%0A%20%20%20%20configure_seaborn()%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23Reinforcement%20Learning%20Chapter%204%3A%20Dynamic%20Programming%0A%0A%20%20%20%20Author%3A%20No%C3%A9%20Lallouet%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Introduction%0A%0A%20%20%20%20Dynamic%20Programming%20refers%20to%20an%20ensemble%20of%20methods%20that%20can%20compute%20optimal%20policies%20for%20a%20Markov%20decision%20process.%20While%20these%20methods%20are%20able%20to%20solve%20problems%20optimally%2C%20they%20require%20a%20complete%20and%20accurate%20model%20of%20the%20environment%2C%20which%20can%20be%20a%20significant%20limitation%20in%20many%20real-world%20applications.%0A%0A%20%20%20%20However%2C%20dynamic%20programming%20methods%20introduce%20us%20to%20several%20important%20concepts%20in%20reinforcement%20learning%2C%20such%20as%20value%20functions%2C%20policies%2C%20and%20the%20Bellman%20equations.%0A%0A%20%20%20%20%23%23%20Problem%20Definition%0A%0A%20%20%20%20We%20are%20using%20the%20Gridworld%20problem%20for%20value%20iteration.%0A%0A%20%20%20%20A%20gridworld%20problem%20is%20a%20simple%20environment%20where%20an%20agent%20navigates%20a%20grid%20to%20reach%20a%20goal%20while%20avoiding%20obstacles.%20Each%20cell%20in%20the%20grid%20represents%20a%20state%2C%20and%20the%20agent%20can%20take%20actions%20to%20move%20in%20different%20directions%20(up%2C%20down%2C%20left%2C%20right).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20Add%20a%20button%20to%20set%20slipery%20variable%20to%20true%20or%20false%0A%20%20%20%20slippery_toggle%20%3D%20mo.ui.switch(label%3D%22Slippery%22%2C%20value%3DFalse)%0A%20%20%20%20grid_size%20%3D%20mo.ui.slider(4%2C%2010%2C%20value%3D8%2C%20step%3D1%2C%20label%3D%22Grid%20Size%22)%0A%20%20%20%20return%20grid_size%2C%20slippery_toggle%0A%0A%0A%40app.cell%0Adef%20_(grid_size%2C%20slippery_toggle)%3A%0A%20%20%20%20env%20%3D%20gym.make(%22FrozenLake-v1%22%2C%20desc%3Dgenerate_random_map(size%3Dgrid_size.value)%2C%20is_slippery%3Dslippery_toggle.value%2C%20render_mode%3D%22rgb_array%22)%0A%20%20%20%20n_states%20%3D%20env.observation_space.n%0A%20%20%20%20n_actions%20%3D%20env.action_space.n%0A%20%20%20%20%23%20Display%20the%20environment%0A%20%20%20%20env.reset()%0A%20%20%20%20return%20(env%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(grid_size%2C%20plot_lake%2C%20slippery_toggle)%3A%0A%20%20%20%20mo.md(rf%22%22%22%0A%20%20%20%20Below%2C%20we%20plot%20an%20image%20of%20the%20gridworld%20environment.%20The%20agent%20must%20navigate%20from%20the%20starting%20point%20to%20the%20goal%20while%20avoiding%20holes.%0A%0A%20%20%20%20If%20the%20%60slippery%60%20property%20is%20activated%2C%20the%20agent%20has%20a%20chance%20to%20slip%20in%20an%20orthogonal%20direction%20while%20taking%20an%20action.%20For%20example%2C%20taking%20the%20action%20%22%E2%86%93%22%20results%20in%20%22%E2%86%93%22%20with%20a%20probability%20%241-%5Cvarepsilon%24%2C%20but%20the%20agent%20may%20slip%20and%20take%20action%20%22%E2%86%90%22%20or%20%22%E2%86%92%22%20with%20probability%20%24%5Cvarepsilon%24.%0A%0A%20%20%20%20%7Bslippery_toggle%7D%0A%0A%20%20%20%20%7Bgrid_size%7D%0A%0A%20%20%20%20%7Bmo.as_html(plot_lake())%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Prediction%20and%20control%0A%0A%20%20%20%20Dynamic%20programming%20algorithms%20leverage%20two%20concepts%3A%0A%20%20%20%20-%20Prediction%3A%20computing%20the%20value%20function%20%24v_%5Cpi%24%20for%20a%20given%20policy%20%24%5Cpi%24%0A%20%20%20%20-%20Control%3A%20finding%20the%20optimal%20policy%20%24%5Cpi%5E*%24%20that%20maximizes%20the%20value%20function.%0A%0A%20%20%20%20We%20will%20use%20the%20following%20algorithms%20for%20prediction%20and%20control%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Policy%20Evaluation%0A%0A%20%20%20%20The%20policy%20evaluation%20algorithm%20computes%20the%20value%20function%20for%20a%20given%20policy.%0A%0A%20%20%20%20Recall%20the%20Bellman%20equation%20for%20the%20value%20function%20%24v_%5Cpi%24%3A%0A%20%20%20%20%24v_%5Cpi(s)%20%3D%20%5Csum_a%20%5Cpi(a%7Cs)%20%5Csum_%7Bs'%2C%20r%7D%20p(s'%2C%20r%20%7C%20s%2C%20a)%20%5Br%20%2B%20%5Cgamma%20v_%5Cpi(s')%5D%24.%0A%0A%20%20%20%20Policy%20evaluation%20transforms%20the%20Bellman%20equation%20into%20an%20update%20rule%20that%20iteratively%20constructs%20a%20sequence%20of%20value%20functions%20that%20converge%20to%20%24v_%5Cpi%24.%20This%20update%20rule%20is%3A%0A%20%20%20%20%24v_%7Bk%2B1%7D(s)%20%3D%20%5Csum_a%20%5Cpi(a%7Cs)%20%5Csum_%7Bs'%2C%20r%7D%20p(s'%2C%20r%20%7C%20s%2C%20a)%20%5Br%20%2B%20%5Cgamma%20v_k(s')%5D%24.%20The%20value%20function%20is%20updated%20for%20all%20states%20in%20the%20MDP.%0A%0A%20%20%20%20Instead%20of%20waiting%20for%20convergence%2C%20it%20is%20common%20to%20stop%20the%20value%20function%20update%20once%20the%20maximum%20change%20in%20the%20value%20function%20across%20all%20states%20is%20below%20a%20certain%20threshold%20%24%5Ctheta%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Policy%20Improvement%0A%0A%20%20%20%20The%20policy%20improvement%20algorithm%20updates%20the%20policy%20based%20on%20the%20current%20value%20function%20%24v_%5Cpi%24%2C%20by%20making%20it%20greedy%20with%20respect%20to%20the%20value%20function.%0A%0A%20%20%20%20For%20all%20states%20in%20the%20Markov%20decision%20process%2C%20the%20policy%20is%20updated%20as%20follows%3A%0A%0A%20%20%20%20%24%24%5Cpi'(s)%20%3D%20%5Carg%5Cmax_a%20%5Csum_%7Bs'%2C%20r%7D%20p(s'%2C%20r%20%7C%20s%2C%20a)%20%5Br%20%2B%20%5Cgamma%20v_%5Cpi(s')%5D%24%24%0A%0A%20%20%20%20The%20Policy%20Improvement%20Theorem%20states%20that%20if%20the%20policy%20%24%5Cpi'%24%20is%20obtained%20from%20%24%5Cpi%24%20by%20making%20it%20greedy%20with%20respect%20to%20%24v_%5Cpi%24%2C%20then%20%24%5Cpi'%20%5Csucceq%20%5Cpi%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Policy%20Iteration%0A%0A%20%20%20%20The%20policy%20iteration%20algorithm%20alternates%20between%20policy%20evaluation%20and%20policy%20improvement%20until%20convergence.%0A%0A%20%20%20%20Here%20is%20an%20outline%20of%20the%20algorithm%3A%0A%0A%20%20%20%20**1.%20Initialization**%0A%0A%20%20%20%20Initialize%20a%20random%20policy%20%24%5Cpi%24%20and%20a%20value%20function%20%24V(s)%24.%0A%0A%20%20%20%20**2.%20Policy%20Evaluation**%0A%0A%20%20%20%20%2B%20Until%20convergence%20under%20a%20threshold%20%24%5Ctheta%24%3A%0A%20%20%20%20%20%20%20%20%2B%20For%20each%20state%20%24s%24%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%20%24V(s)%20%5Cleftarrow%20%5Csum_%7Bs'%2C%20r%7D%20p(s'%2C%20r%20%7C%20s%2C%20%5Cpi(s))%20%5Br%20%2B%20%5Cgamma%20V(s')%5D%24%0A%0A%20%20%20%20**3.%20Policy%20Improvement**%0A%0A%20%20%20%20%2B%20For%20each%20state%20%24s%24%3A%0A%20%20%20%20%20%20%20%20%2B%20%24%5Cpi(s)%20%5Cleftarrow%20%5Carg%5Cmax_a%20%5Csum_%7Bs'%2C%20r%7D%20p(s'%2C%20r%20%7C%20s%2C%20a)%20%5Br%20%2B%20%5Cgamma%20V(s')%5D%24%0A%20%20%20%20%2B%20If%20the%20policy%20is%20stable%2C%20stop.%20Otherwise%2C%20go%20to%20step%202.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.class_definition(hide_code%3DTrue)%0Aclass%20Algorithm%3A%20%0A%0A%20%20%20%20def%20__init__(self%2C%20env%2C%20gamma%3D0.9%2C%20epsilon%3D1e-3%2C%20n_iter%3D1000)%3A%0A%20%20%20%20%20%20%20%20self.env%20%3D%20env%0A%20%20%20%20%20%20%20%20self.gamma%20%3D%20gamma%0A%20%20%20%20%20%20%20%20self.epsilon%20%3D%20epsilon%0A%20%20%20%20%20%20%20%20self.n_iter%20%3D%20n_iter%0A%20%20%20%20%20%20%20%20self.n_states%20%3D%20env.observation_space.n%0A%20%20%20%20%20%20%20%20self.V%20%3D%20np.zeros(env.observation_space.n)%0A%0A%0A%40app.class_definition%0Aclass%20PolicyIteration(Algorithm)%3A%0A%0A%20%20%20%20def%20__init__(self%2C%20env%2C%20gamma%3D0.9%2C%20epsilon%3D1e-3%2C%20n_iter%3D1000)%3A%0A%20%20%20%20%20%20%20%20super().__init__(env%2C%20gamma%2C%20epsilon%2C%20n_iter)%0A%20%20%20%20%20%20%20%20self.V%20%3D%20np.zeros(env.observation_space.n)%0A%20%20%20%20%20%20%20%20self.policy%20%3D%20np.random.choice(env.action_space.n%2C%20size%3Dself.n_states)%0A%20%20%20%20%20%20%20%20self.v_archive%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20self.pi_archive%20%3D%20%5B%5D%0A%0A%20%20%20%20def%20policy_evaluation(self)%3A%0A%20%20%20%20%20%20%20%20while%20True%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20V_old%20%3D%20self.V.copy()%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20state%20in%20range(self.n_states)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20action%20%3D%20self.policy%5Bstate%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20v_value%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20prob%2C%20next_state%2C%20reward%2C%20done%20in%20self.env.unwrapped.P%5Bstate%5D%5Baction%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20v_value%20%2B%3D%20prob%20*%20(reward%20%2B%20self.gamma%20*%20self.V%5Bnext_state%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.V%5Bstate%5D%20%3D%20v_value%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20np.max(np.abs(self.V%20-%20V_old))%20%3C%20self.epsilon%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20break%0A%0A%20%20%20%20def%20policy_improvement(self)%3A%0A%20%20%20%20%20%20%20%20policy_stable%20%3D%20True%0A%20%20%20%20%20%20%20%20for%20state%20in%20range(self.n_states)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20old_action%20%3D%20self.policy%5Bstate%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20q_values%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20actions%20%3D%20range(self.env.action_space.n)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20action%20in%20actions%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20prob%2C%20next_state%2C%20reward%2C%20done%20in%20self.env.unwrapped.P%5Bstate%5D%5Baction%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%2B%3D%20prob%20*%20(reward%20%2B%20self.gamma%20*%20self.V%5Bnext_state%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_values.append(q_value)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy%5Bstate%5D%20%3D%20np.argmax(q_values)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20old_action%20!%3D%20self.policy%5Bstate%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_stable%20%3D%20False%0A%20%20%20%20%20%20%20%20return%20policy_stable%0A%0A%20%20%20%20def%20run(self)%3A%0A%20%20%20%20%20%20%20%20for%20i%20in%20range(self.n_iter)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Archive%20the%20value%20function%20and%20policy%0A%20%20%20%20%20%20%20%20%20%20%20%20self.v_archive.append(self.V.copy())%0A%20%20%20%20%20%20%20%20%20%20%20%20self.pi_archive.append(self.policy.copy())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Policy%20Evaluation%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_evaluation()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Policy%20Improvement%0A%20%20%20%20%20%20%20%20%20%20%20%20policy_stable%20%3D%20self.policy_improvement()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20policy_stable%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(%22Policy%20converged%20at%20iteration%3A%22%2C%20i)%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%20self.V%2C%20self.policy%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Let's%20run%20the%20policy%20iteration%20algorithm.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(env)%3A%0A%20%20%20%20policy_iteration%20%3D%20PolicyIteration(env%2C%20gamma%3D0.9%2C%20epsilon%3D1e-8)%0A%20%20%20%20v_policy%2C%20pi_policy%20%3D%20policy_iteration.run()%0A%20%20%20%20return%20(policy_iteration%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(algorithm%2C%20plot_all%2C%20slider)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20Here%2C%20we%20plot%20the%20value%20function%20and%20the%20associated%20policy%20for%20all%20the%20states%20in%20the%20Gridworld.%0A%0A%20%20%20%20The%20%7Bslider%7D%20allows%20to%20view%20the%20evolution%20of%20the%20value%20function%20over%20iterations.%0A%20%20%20%20%7Bmo.as_html(plot_all(algorithm%2C%20slider.value))%7D%0A%20%20%20%20We%20can%20also%20visualize%20the%20optimal%20policy%20derived%20from%20the%20value%20function.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Value%20Iteration%0A%0A%20%20%20%20The%20policy%20iteration%20algorithm%20requires%20policy%20evaluation%20updates%20at%20every%20iteration%2C%20which%20may%20be%20computationally%20expensive.%20In%20order%20to%20accelerate%20computation%2C%20the%20value%20iteration%20algorithm%20combines%20policy%20improvement%20with%20a%20truncated%20step%20of%20policy%20evaluation.%20In%20the%20limit%2C%20value%20iteration%20correctly%20converges%20to%20the%20optimal%20policy.%0A%0A%20%20%20%20The%20value%20iteration%20algorithm%20iteratively%20updates%20the%20value%20function%20until%20convergence.%0A%0A%20%20%20%20The%20update%20rule%20combines%20policy%20improvement%20with%20a%20single%20sweep%20of%20policy%20evaluation%3A%0A%20%20%20%20%24v_%7Bk%2B1%7D(s)%20%3D%20%5Cmax_a%20%5Csum_%7Bs'%2C%20r%7D%20P(s'%2Cr%7Cs%2Ca)%20%5Br%20%2B%20%5Cgamma%20v_k(s')%5D%24%0A%0A%20%20%20%20The%20value%20iteration%20algorithm%20is%20the%20following%3A%0A%0A%20%20%20%20**1.%20Initialization**%0A%0A%20%20%20%20-%20Initialize%20a%20random%20policy%20%24%5Cpi%24%20and%20a%20value%20function%20%24V(s)%24.%0A%0A%20%20%20%20**2.%20Value%20Update**%0A%0A%20%20%20%20-%20Loop%20until%20convergence%20of%20%24%5CDelta%24%20under%20a%20threshold%20%24%5Ctheta%24%3A%0A%20%20%20%20%20%20%20%20-%20For%20each%20state%20%24s%24%3A%20%24v_%7Bk%2B1%7D(s)%20%3D%20%5Cmax_a%20%5Csum_%7Bs'%2C%20r%7D%20P(s'%2Cr%7Cs%2Ca)%20%5Br%20%2B%20%5Cgamma%20v_k(s')%5D%24%0A%20%20%20%20%20%20%20%20-%20%24%5CDelta%20%5Cleftarrow%24%20max%24_s%20%7Cv_%7Bk%2B1%7D(s)%20-%20v_k(s)%7C%24%0A%0A%20%20%20%20**3.%20Policy%20Computation**%0A%0A%20%20%20%20-%20For%20each%20state%20%24s%24%3A%20%24%5Cpi(s)%20%5Cleftarrow%20%5Carg%5Cmax_a%20%5Csum_%7Bs'%2C%20r%7D%20P(s'%2Cr%7Cs%2Ca)%20%5Br%20%2B%20%5Cgamma%20v_k(s')%5D%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.class_definition%0Aclass%20ValueIteration(Algorithm)%3A%0A%0A%20%20%20%20def%20__init__(self%2C%20env%2C%20gamma%3D0.9%2C%20epsilon%3D1e-3%2C%20n_iter%3D1000)%3A%0A%20%20%20%20%20%20%20%20super().__init__(env%2C%20gamma%2C%20epsilon%2C%20n_iter)%0A%20%20%20%20%20%20%20%20self.v_archive%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20self.pi_archive%20%3D%20%5B%5D%0A%0A%20%20%20%20def%20update(self)%3A%0A%20%20%20%20%20%20%20%20for%20state%20in%20range(self.n_states)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20q_values%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20actions%20%3D%20range(self.env.action_space.n)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20action%20in%20actions%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20prob%2C%20next_state%2C%20reward%2C%20done%20in%20self.env.unwrapped.P%5Bstate%5D%5Baction%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%2B%3D%20prob%20*%20(reward%20%2B%20self.gamma%20*%20self.V%5Bnext_state%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_values.append(q_value)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.V%5Bstate%5D%20%3D%20max(q_values)%0A%0A%20%20%20%20def%20return_policy(self)%3A%0A%20%20%20%20%20%20%20%20policy%20%3D%20np.zeros(self.n_states%2C%20dtype%3Dint)%0A%20%20%20%20%20%20%20%20for%20state%20in%20range(self.n_states)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20q_values%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20actions%20%3D%20range(self.env.action_space.n)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20action%20in%20actions%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20prob%2C%20next_state%2C%20reward%2C%20done%20in%20self.env.unwrapped.P%5Bstate%5D%5Baction%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_value%20%2B%3D%20prob%20*%20(reward%20%2B%20self.gamma%20*%20self.V%5Bnext_state%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20q_values.append(q_value)%0A%20%20%20%20%20%20%20%20%20%20%20%20policy%5Bstate%5D%20%3D%20np.argmax(q_values)%0A%20%20%20%20%20%20%20%20return%20policy%0A%0A%20%20%20%20def%20run(self)%3A%0A%20%20%20%20%20%20%20%20for%20i%20in%20range(self.n_iter)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20V_old%20%3D%20self.V.copy()%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Archive%20the%20value%20function%20and%20policy%0A%20%20%20%20%20%20%20%20%20%20%20%20self.v_archive.append(self.V.copy())%0A%20%20%20%20%20%20%20%20%20%20%20%20self.pi_archive.append(self.return_policy())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Update%20the%20value%20function%0A%20%20%20%20%20%20%20%20%20%20%20%20self.update()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20np.max(np.abs(self.V%20-%20V_old))%20%3C%20self.epsilon%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(%22Convergence%20reached%20at%20iteration%3A%22%2C%20i)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20break%0A%20%20%20%20%20%20%20%20return%20self.V%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Let's%20run%20the%20algorithm.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(env)%3A%0A%20%20%20%20algorithm%20%3D%20ValueIteration(env%2C%20gamma%3D0.9%2C%20epsilon%3D1e-3)%0A%20%20%20%20v%20%3D%20algorithm.run()%0A%20%20%20%20return%20(algorithm%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(plot_all%2C%20policy_iteration%2C%20slider_pi)%3A%0A%20%20%20%20mo.md(rf%22%22%22%0A%20%20%20%20Here%2C%20we%20plot%20the%20value%20function%20and%20the%20associated%20policy%20for%20all%20the%20states%20in%20the%20Gridworld.%0A%0A%20%20%20%20The%20%7Bslider_pi%7D%20allows%20to%20view%20the%20evolution%20of%20the%20value%20function%20%24v_%5Cpi%24%20and%20the%20associated%20policy%20%24%5Cpi%24%20over%20the%20course%20of%20the%20convergence%20of%20the%20algorithm.%0A%20%20%20%20%7Bmo.as_html(plot_all(policy_iteration%2C%20slider_pi.value))%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(algorithm)%3A%0A%20%20%20%20archive_length%20%3D%20len(algorithm.v_archive)%0A%20%20%20%20slider%20%3D%20mo.ui.slider(0%2C%20archive_length-1%2C%20value%3Darchive_length-1%2C%20step%3D1%2C%20label%3D%22iteration%20slider%22)%0A%20%20%20%20return%20(slider%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(policy_iteration)%3A%0A%20%20%20%20archive_length_pi%20%3D%20len(policy_iteration.v_archive)%0A%20%20%20%20slider_pi%20%3D%20mo.ui.slider(0%2C%20archive_length_pi-1%2C%20value%3Darchive_length_pi-1%2C%20step%3D1%2C%20label%3D%22iteration%20slider%22)%0A%20%20%20%20return%20(slider_pi%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(env)%3A%0A%20%20%20%20from%20matplotlib.colors%20import%20LinearSegmentedColormap%0A%0A%20%20%20%20def%20slice_cmap(cmap%2C%20start%3D0.25%2C%20end%3D0.75%2C%20name%3DNone%2C%20N%3D256)%3A%0A%20%20%20%20%20%20%20%20if%20name%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20name%20%3D%20f%22%7Bcmap.name%7D_slice_%7Bstart%3A.2f%7D_%7Bend%3A.2f%7D%22%0A%20%20%20%20%20%20%20%20x%20%3D%20np.linspace(start%2C%20end%2C%20N)%0A%20%20%20%20%20%20%20%20colors%20%3D%20cmap(x)%0A%20%20%20%20%20%20%20%20return%20LinearSegmentedColormap.from_list(name%2C%20colors%2C%20N%3DN)%0A%0A%20%20%20%20%40functools.cache%0A%20%20%20%20def%20plot_value_function(algorithm%2C%20iteration)%3A%0A%20%20%20%20%20%20%20%20v_iter%20%3D%20algorithm.v_archive%5Biteration%5D%0A%20%20%20%20%20%20%20%20v_array_iter%20%3D%20v_iter.reshape((algorithm.env.unwrapped.nrow%2C%20algorithm.env.unwrapped.ncol))%0A%20%20%20%20%20%20%20%20return%20v_array_iter%0A%0A%20%20%20%20%40functools.cache%0A%20%20%20%20def%20plot_policy(algorithm%2C%20iteration)%3A%0A%20%20%20%20%20%20%20%20action_mapping%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%200%3A%20'%E2%86%90'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%201%3A%20'%E2%86%93'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%202%3A%20'%E2%86%92'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%203%3A%20'%E2%86%91'%2C%0A%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20policy%20%3D%20algorithm.pi_archive%5Biteration%5D%0A%20%20%20%20%20%20%20%20policy_symbols%20%3D%20%5Baction_mapping%5Baction%5D%20for%20action%20in%20policy%5D%0A%20%20%20%20%20%20%20%20policy_array%20%3D%20np.array(policy_symbols).reshape((algorithm.env.unwrapped.nrow%2C%20algorithm.env.unwrapped.ncol))%0A%20%20%20%20%20%20%20%20%23Where%20there%20are%20obstacles%2C%20set%20policy%20to%20empty%20string%0A%20%20%20%20%20%20%20%20%23%20Get%20location%20of%20obstacles%20in%20env%0A%20%20%20%20%20%20%20%20for%20row%20in%20range(env.unwrapped.nrow)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20col%20in%20range(env.unwrapped.ncol)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state%20%3D%20row%20*%20env.unwrapped.ncol%20%2B%20col%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20env.unwrapped.desc%5Brow%2C%20col%5D%20%3D%3D%20b'H'%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_array%5Brow%2C%20col%5D%20%3D%20'H'%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20env.unwrapped.desc%5Brow%2C%20col%5D%20%3D%3D%20b'G'%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_array%5Brow%2C%20col%5D%20%3D%20'G'%0A%20%20%20%20%20%20%20%20return%20policy%2C%20policy_array%0A%0A%0A%20%20%20%20%40functools.cache%0A%20%20%20%20def%20plot_all(algorithm%2C%20iteration)%3A%0A%20%20%20%20%20%20%20%20n_row%20%3D%20algorithm.env.unwrapped.nrow%0A%20%20%20%20%20%20%20%20n_col%20%3D%20algorithm.env.unwrapped.ncol%0A%20%20%20%20%20%20%20%20f_main%2C%20(ax_value%2C%20ax_policy)%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12%2C%205))%0A%20%20%20%20%20%20%20%20cbar_ax%20%3D%20f_main.add_axes(%5B0.48%2C%200.1%2C%200.01%2C%200.78%5D)%0A%20%20%20%20%20%20%20%20v%20%3D%20plot_value_function(algorithm%2C%20iteration)%0A%20%20%20%20%20%20%20%20p%2C%20p_sym%20%3D%20plot_policy(algorithm%2C%20iteration)%0A%20%20%20%20%20%20%20%20sns.heatmap(v%2C%20annot%3DTrue%2C%20cmap%3D%22YlGnBu%22%2C%20cbar%3DTrue%2C%20ax%3Dax_value%2C%20square%3DTrue%2C%20cbar_ax%3Dcbar_ax%2C%20fmt%3D%22.2f%22)%0A%20%20%20%20%20%20%20%20sns.heatmap(p.reshape((n_row%2C%20n_col))%2C%20annot%3Dp_sym%2C%20fmt%3D''%2C%20cmap%3Dslice_cmap(plt.get_cmap('Blues')%2C%200.1%2C%200.5)%2C%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cbar%3DFalse%2C%20ax%3Dax_policy%2C%20square%3DTrue)%0A%20%20%20%20%20%20%20%20mask_hole%20%3D%20np.zeros_like(p.reshape((n_row%2C%20n_col))%2C%20dtype%3Dbool)%0A%20%20%20%20%20%20%20%20mask_goal%20%3D%20np.zeros_like(p.reshape((n_row%2C%20n_col))%2C%20dtype%3Dbool)%0A%20%20%20%20%20%20%20%20for%20row%20in%20range(n_row)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20col%20in%20range(n_col)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20env.unwrapped.desc%5Brow%2C%20col%5D%20%3D%3D%20b'H'%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mask_hole%5Brow%2C%20col%5D%20%3D%20True%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20env.unwrapped.desc%5Brow%2C%20col%5D%20%3D%3D%20b'G'%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mask_goal%5Brow%2C%20col%5D%20%3D%20True%0A%0A%20%20%20%20%20%20%20%20for%20axis%2C%20symbols%2C%20format%20in%20zip(%5Bax_value%2C%20ax_policy%5D%2C%20%5Bv%2C%20p_sym%5D%2C%20%5B'.2f'%2C%20''%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20axis.set_xticks(%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20axis.set_yticks(%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20sns.heatmap(np.ones((n_row%2C%20n_col))%2C%20mask%3D~mask_hole%2C%20annot%3Dsymbols%2C%20fmt%3Dformat%2C%20cmap%3D%5B%22%234a5267%22%5D%2C%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cbar%3DFalse%2C%20ax%3Daxis%2C%20square%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20sns.heatmap(np.ones((n_row%2C%20n_col))%2C%20mask%3D~mask_goal%2C%20annot%3Dsymbols%2C%20fmt%3Dformat%2C%20cmap%3D%5B%22%23ffcd00%22%5D%2C%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cbar%3DFalse%2C%20ax%3Daxis%2C%20square%3DTrue)%0A%0A%20%20%20%20%20%20%20%20ax_value.set_title(f%22Value%20Function%20at%20iteration%20%7Biteration%7D%22)%0A%20%20%20%20%20%20%20%20ax_policy.set_title(f%22Policy%20at%20iteration%20%7Biteration%7D%22)%0A%20%20%20%20%20%20%20%20return%20plt.gca()%0A%20%20%20%20return%20(plot_all%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(env%2C%20slippery_toggle)%3A%0A%20%20%20%20%40functools.cache%0A%20%20%20%20def%20plot_lake()%3A%0A%20%20%20%20%20%20%20%20try%3A%20%0A%20%20%20%20%20%20%20%20%20%20%20%20lake%20%3D%20env.render()%0A%20%20%20%20%20%20%20%20%20%20%20%20f%2C%20ax%20%3D%20plt.subplots()%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.imshow(lake)%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.set_title(f%22Frozen%20Lake%20Environment%20(slippery%3D%7Bslippery_toggle.value%7D)%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.set_xticks(%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.set_yticks(%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20plt.gca()%0A%20%20%20%20%20%20%20%20except%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20lake%20%3D%20env.unwrapped._render_text()%0A%20%20%20%20%20%20%20%20%20%20%20%20print(lake)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%0A%20%20%20%20return%20(plot_lake%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%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
d08ff9b449c974b3e900dc1cb0bc0e98