Create training.py

training.py

@@ -0,0 +1,106 @@
+import torch
+from agent_class import ParameterisedPolicy
+
+def create_cum_rewards(rewards, discount=DISCOUNT):
+    new_rews = [0]
+    for el in rewards[::-1]:
+        val = el + discount * new_rews[-1]
+        new_rews.append(val)
+    return torch.tensor(new_rews[1:][::-1], dtype=torch.float32)
+
+
+def play_game(env, model, n_steps=500, render=False):
+    observation = env.reset()
+
+    rewards, logits = [], []
+#     for _ in range(n_steps):
+    while True:
+        if render:
+            env.render()
+            
+        (mus, sigmas) = model(torch.tensor(observation, dtype=torch.float32))
+        
+        m = torch.distributions.normal.Normal(mus, sigmas)
+        action = m.sample()
+        logit = m.log_prob(action)
+        observation, reward, done, info = env.step(action.detach().numpy())
+
+        rewards.append(reward)
+        logits.append(m.log_prob(action).sum())
+        
+        if done:
+            break
+    env.close()
+    
+    return rewards, logits
+
+def draw_gradients_rewards(model, rewards, ep_lengths, ave_over_steps):
+    
+    fig, axs = plt.subplot_mosaic([['1', '1', '2', '2'], ['3', '4', '5', '6']],
+                                  constrained_layout=False, figsize=(20, 9))
+
+    axs['1'].plot(np.array(rewards[:ave_over_steps*(len(rewards)//ave_over_steps)])\
+                    .reshape(-1, ave_over_steps).mean(axis=-1))
+    axs['1'].set_title('Sum rewards per episode')
+    
+    axs['1'].hlines(200, 0, len(rewards)/ave_over_steps, colors='red')
+    axs['1'].hlines(150, 0, len(rewards)/ave_over_steps, colors='orange')
+    axs['1'].hlines(0, 0, len(rewards)/ave_over_steps, colors='green')
+    
+    axs['2'].plot(np.array(ep_lengths[:ave_over_steps*(len(ep_lengths)//ave_over_steps)])\
+                    .reshape(-1, ave_over_steps).mean(axis=-1))
+    axs['2'].set_title('Episode length')
+    
+    axs['3'].hist(model.lin_1.weight.grad.flatten().detach().numpy(), bins=50); 
+    axs['3'].set_xlabel('Grads in dense layer 1')
+
+    axs['4'].hist(model.lin_2.weight.grad.flatten().detach().numpy(), bins=50); 
+    axs['4'].set_xlabel('Grads in dense layer 2')
+
+    axs['5'].hist(model.lin_3.weight.grad.flatten().detach().numpy(), bins=50); 
+    axs['5'].set_xlabel('Grads in dense layer 3')
+    
+    axs['6'].hist(model.lin_4.weight.grad.flatten().detach().numpy(), bins=50); 
+    axs['6'].set_xlabel('Grads in dense layer 4')
+    
+model = ParameterisedPolicy()
+opt = torch.optim.Adam(model.parameters(), lr=0.0008)
+lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=4000, gamma=0.7)
+rews, ep_lengths = [], []
+
+last_max_score = 50
+env = gym.make(env_name)
+
+for _ in range(int(10e3)):
+    rewards, logits = play_game(env, model, render=False)
+    
+    cum_rewards = create_cum_rewards(rewards, discount=DISCOUNT)
+    stacked_logits = torch.stack(logits).flatten()
+    
+    loss = -(stacked_logits * cum_rewards).mean()
+    
+    rews.append(np.sum(rewards))
+    ep_lengths.append(len(rewards))
+
+    opt.zero_grad()
+    loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 50)
+    opt.step()
+    lr_scheduler.step()
+    
+    if _%40 == 0:
+        if _ > 1:
+            clear_output()
+            draw_gradients_rewards(model, rewards=rews, 
+                                   ep_lengths=ep_lengths, ave_over_steps=40)
+            plt.show()
+            
+    if len(rews) > 40:
+        agg_rews = np.array(rews[-40*(len(rews)//40):])\
+                        .reshape(-1, 40).mean(axis=-1) 
+        if (agg_rews[-1] > last_max_score):
+            last_max_score = agg_rews[-1]
+            print('NEW BEST MODEL, STEP:', _, 'SCORE: ', last_max_score)
+            save_path = f'best_models/best_reinforce_lunar_lander_cont_model_{round(last_max_score,3)}.pt'
+            torch.save(model, save_path)
+