New branch

Add Tensorboard visualisation to quad training file

Author: ArzhelaR

model_RL/PPO_model_pers.py

@@ -199,14 +199,16 @@ def train(self, dataset):
                     p.requires_grad = True
                 start = stop + 1
 
-    def learn(self):
+    def learn(self, writer):
         """
         Train the PPO mesh_model
         :return: the actor policy, training rewards, training wins, len of episodes
         """
         rewards = []
         wins = []
         len_ep = []
+        global_step = 0
+        nb_episodes = 0
 
         try:
             for iteration in tqdm(range(self.nb_iterations)):
@@ -217,6 +219,8 @@ def learn(self):
                     next_obs, info = self.env.reset()
                     trajectory = []
                     ep_reward = 0
+                    ep_mesh_reward = 0
+                    ideal_reward = info["mesh_ideal_rewards"]
                     done = False
                     step = 0
                     while step < 40:
@@ -229,6 +233,7 @@ def learn(self):
                         gym_action = [action[2],int(action[0]/3)]
                         next_obs, reward, terminated, truncated, info = self.env.step(gym_action)
                         ep_reward += reward
+                        ep_mesh_reward += info["mesh_reward"]
                         if terminated:
                             if truncated:
                                 wins.append(0)
@@ -245,6 +250,10 @@ def learn(self):
                         rollouts.append(trajectory)
                         dataset.extend(trajectory)
                         len_ep.append(len(trajectory))
+                    nb_episodes += 1
+                    writer.add_scalar("episode_reward", ep_reward, nb_episodes)
+                    writer.add_scalar("normalized return", (ep_reward/ideal_reward), nb_episodes)
+                    writer.add_scalar("len_episodes", len(trajectory), nb_episodes)
 
                 self.train(dataset)
 

training/train_quadmesh.py

@@ -10,20 +10,26 @@
 from model_RL.PPO_model_pers import PPO
 
 import gymnasium as gym
+from torch.utils.tensorboard import SummaryWriter
+import random
+import torch
+import numpy as np
+import time
+import wandb
 import json
+import os
 
+if __name__ == '__main__':
 
-def train():
-    mesh_size = 30
-    lr = 0.0001
-    gamma = 0.9
-
-    #dataset = [random_mesh() for _ in range(9)]
-    #plot_dataset(dataset)
-
+    with open("model_RL/parameters/ppo_config.json", "r") as f:
+        ppo_config = json.load(f)
     with open("environment/environment_config.json", "r") as f:
         env_config = json.load(f)
 
+    # Create log dir
+    log_dir = ppo_config["tensorboard_log"]
+    os.makedirs(log_dir, exist_ok=True)
+
     # Create the environment
     env = gym.make(
         env_config["env_name"],
@@ -35,16 +41,40 @@ def train():
         with_degree_obs=env_config["with_degree_observation"]
     )
 
-    model = PPO(env, lr, gamma, nb_iterations=15, nb_episodes_per_iteration=100, nb_epochs=5, batch_size=8)
-    actor, rewards, wins, steps = model.learn()
-    if rewards is not None:
-        plot_training_results(rewards, wins, steps)
+    model = PPO(
+        env=env,
+        lr=ppo_config["learning_rate"],
+        gamma=ppo_config["gamma"],
+        nb_iterations=20,
+        nb_episodes_per_iteration=100,
+        nb_epochs=5,
+        batch_size=8
+    )
 
-"""
-    # torch.save(actor.state_dict(), 'policy_saved/actor_network.pth')
-    avg_steps, avg_wins, avg_rewards, final_meshes = testPolicy(actor, 5, dataset, 60)
+    run_name = f"{env_config['env_name']}__{1}__{int(time.time())}"
+    # Create log dir
+    log_dir = ppo_config["tensorboard_log"]
+    os.makedirs(log_dir, exist_ok=True)
 
+    # SEEDING
+    seed = 1
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+    writer = SummaryWriter(f"results/runs/{run_name}")
+    writer.add_text(
+        "Environment config",
+        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in env_config.items()])),
+    )
+    writer.add_text(
+        "PPO config",
+        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in ppo_config.items()])),
+    )
+
+    actor, rewards, wins, steps = model.learn(writer)
+    writer.close()
     if rewards is not None:
-        plot_test_results(avg_rewards, avg_wins, avg_steps, avg_rewards)
-    plot_dataset(final_meshes)
-"""
+        plot_training_results(rewards, wins, steps)
+    # torch.save(actor.state_dict(), 'policy_saved/actor_network.pth')