add gitignore and training loop class

.gitignore

@@ -0,0 +1 @@
+.DS_Store

ddpg.py

@@ -176,20 +176,23 @@ class Agent(object):
         self.target_critic = CriticNetwork(beta, input_dims, layer1_size, layer2_size, n_actions=n_actions, name="target_critic")
 
         self.noise = OUActionNoise(mu=np.zeros(n_actions))
+
+        self.attributions = None
+        self.ig = None
         
         self.update_network_parameters(tau=1)
 
-    def choose_action(self, observation, attribution : IntegratedGradients = None, baseline : np.ndarray=None):
+    def choose_action(self, observation, baseline : T.Tensor=None):
         self.actor.eval()
         observation = T.tensor(observation, dtype=T.float).to(self.actor.device)
         print(f"Observation: {observation.shape=}")
         mu = self.actor(observation).to(self.actor.device)
 
-        if attribution is not None:
-            if baseline is None:
-                baseline = T.zeros(observation.shape)
-            attributions = attribution.attribute((observation), baselines=baseline, target=0)
-            print('Attributions:', attributions)
+        # if attribution is not None:
+        #     if baseline is None:
+        #         baseline = T.zeros(observation.shape)
+        #     attributions = attribution.attribute((observation), baselines=baseline, target=0)
+        #     print('Attributions:', attributions)
 
 
         mu_prime = mu + T.tensor(self.noise(), dtype=T.float).to(self.actor.device)

main.py

@@ -7,7 +7,8 @@ import argparse
 from train import TrainingLoop
 from captum.attr import (IntegratedGradients, LayerConductance, NeuronAttribution)
 
-training_loop = TrainingLoop()
+training_loop = TrainingLoop(env_spec="LunarLander-v2", continuous=True, gravity=-10, render_mode=None)
+training_loop.create_agent()
 
 parser = argparse.ArgumentParser(description="Choose a function to run.")
 parser.add_argument("function", choices=["train", "load-trained", "attribute"], help="The function to run.")

train.py

@@ -3,69 +3,75 @@ import gymnasium as gym
 import numpy as np
 import matplotlib.pyplot as plt
 import torch
-import argparse
 from captum.attr import (IntegratedGradients)
 
 
 class TrainingLoop:
-    def __init__(self):
-        pass 
+    def __init__(self, env_spec, output_path='./output/', seed=0, **kwargs):
+        assert env_spec in gym.envs.registry.keys()
+        defaults = {
+            "continuous": True,
+            "gravity": -10.0,
+            "render_mode": None
+        }
 
-    def train(self):
-        env = gym.make(
-            "LunarLander-v2",
-            continuous = True,
-            gravity = -10.0,
-            render_mode = None
+        self.env = gym.make(
+            env_spec,
+            **defaults
         )
 
-        agent = Agent(alpha=0.000025, beta=0.00025, input_dims=[8], tau=0.001, env=env, batch_size=64, layer1_size=400, layer2_size=300, n_actions=4)
-        agent.load_models()
+        torch.manual_seed(seed)
+
+        self.agent = None
+        self.output_path = output_path
+
+    # TODO: spec-to-hyperparameters look-up
+    def create_agent(self, alpha=0.000025, beta=0.00025, input_dims=[8], tau=0.001, batch_size=64, layer1_size=400, layer2_size=300, n_actions=4):
+        self.agent = Agent(alpha=alpha, beta=beta, input_dims=input_dims, tau=tau, env=self.env, batch_size=batch_size, layer1_size=layer1_size, layer2_size=layer2_size, n_actions=n_actions)
+
+    def train(self):
+        assert self.agent is not None
+        
+        self.agent.load_models()
 
-        np.random.seed(0)
         score_history = []
 
         for i in range(1000):
             done = False
             score = 0
-            obs, _ = env.reset()
+            obs, _ = self.env.reset()
             while not done:
-                act = agent.choose_action(obs)
-                new_state, reward, terminated, truncated, info = env.step(act)
+                act = self.agent.choose_action(obs)
+                new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
-                agent.remember(obs, act, reward, new_state, int(done))
-                agent.learn()
+                self.agent.remember(obs, act, reward, new_state, int(done))
+                self.agent.learn()
                 score += reward
                 obs = new_state
 
             score_history.append(score)
             print("episode", i, "score %.2f" % score, "100 game average %.2f" % np.mean(score_history[-100:]))
             if i % 25 == 0:
-                agent.save_models()
+                self.agent.save_models()
+
+        self.env.close()
 
 
     def load_trained(self):
-        env = gym.make(
-            "LunarLanderContinuous-v2",
-            render_mode = "human"
-        )
+        assert self.agent is not None
 
-        agent = Agent(alpha=0.000025, beta=0.00025, input_dims=[8], tau=0.001, env=env, batch_size=64, layer1_size=400, layer2_size=300, n_actions=4)
-        agent.load_models()
+        self.agent.load_models()
 
-        np.random.seed(0)
         score_history = []
 
         for i in range(50):
             done = False
             score = 0
-            obs, _ = env.reset()
+            obs, _ = self.env.reset()
             
-
-
             while not done:
-                act = agent.choose_action(obs)
-                new_state, reward, terminated, truncated, info = env.step(act)
+                act = self.agent.choose_action(obs)
+                new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
                 score += reward
                 obs = new_state
@@ -73,42 +79,42 @@ class TrainingLoop:
             score_history.append(score)
             print("episode", i, "score %.2f" % score, "100 game average %.2f" % np.mean(score_history[-100:]))
 
+        self.env.close()
+
     # Model Explainability
 
     from captum.attr import (IntegratedGradients)
 
     def _collect_running_baseline_average(self, num_iterations: int) -> torch.Tensor:
-        env = gym.make(
-            "LunarLanderContinuous-v2",
-            render_mode = None
-        )
-
-        agent = Agent(alpha=0.000025, beta=0.00025, input_dims=[8], tau=0.001, env=env, batch_size=64, layer1_size=400, layer2_size=300, n_actions=4)
-        agent.load_models()
+        assert self.agent is not None
 
-        torch.manual_seed(0)
+        self.agent.load_models()
 
         sum_obs = torch.zeros(8)
 
         for i in range(num_iterations):
             done = False
             score = 0
-            obs, _ = env.reset()
+            obs, _ = self.env.reset()
             
             sum_obs += obs
             print(f"Baseline on interation #{i}: {obs}")
 
             while not done:
-                act = agent.choose_action(obs, attribution=None, baseline=None)
-                new_state, reward, terminated, truncated, info = env.step(act)
+                act = self.agent.choose_action(obs, attribution=None, baseline=None)
+                new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
                 score += reward
                 obs = new_state
 
+        self.env.close()
+
         return sum_obs / num_iterations
 
 
     def explain_trained(self, option: str, num_iterations :int = 10) -> None:
+        assert self.agent is not None
+        
         baseline_options = {
             "1": torch.zeros(8),
             "2": self._collect_running_baseline_average(num_iterations), 
@@ -116,27 +122,20 @@ class TrainingLoop:
 
         baseline = baseline_options[option]
 
-        env = gym.make(
-            "LunarLanderContinuous-v2",
-            render_mode = "human"
-        )
+        self.agent.load_models()
 
-        agent = Agent(alpha=0.000025, beta=0.00025, input_dims=[8], tau=0.001, env=env, batch_size=64, layer1_size=400, layer2_size=300, n_actions=4)
+        ig = IntegratedGradients(self.agent.actor)
+        self.agent.ig = ig
 
-        agent.load_models()
-
-        ig = IntegratedGradients(agent.actor)
-
-        np.random.seed(0)
         score_history = []
 
         for i in range(50):
             done = False
             score = 0
-            obs, _ = env.reset()
+            obs, _ = self.env.reset()
             while not done:
-                act = agent.choose_action(obs, attribution=ig, baseline=baseline)
-                new_state, reward, terminated, truncated, info = env.step(act)
+                act = self.agent.choose_action(obs, baseline=baseline)
+                new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
                 score += reward
                 obs = new_state
@@ -144,3 +143,7 @@ class TrainingLoop:
             score_history.append(score)
             print("episode", i, "score %.2f" % score, "100 game average %.2f" % np.mean(score_history[-100:]))
 
+        self.env.close()
+
+        return self.agent.attributions
+