working attribution mechanism

ddpg.py

@@ -143,7 +143,12 @@ class ActorNetwork(nn.Module):
         self.to(self.device)
 
     def forward(self, state):
-        print(f"State in forward function: {state.shape=}")
+        
+        try:
+            assert state.shape == T.Size([8])
+        except AssertionError:
+            raise Exception(f"Wrong shape {state.shape=}")
+        
         x = self.fc1(state)
         x = self.bn1(x)
         x = F.relu(x)
@@ -178,21 +183,19 @@ class Agent(object):
         self.noise = OUActionNoise(mu=np.zeros(n_actions))
 
         self.attributions = None
-        self.ig = None
+        self.ig : IntegratedGradients = None
         
         self.update_network_parameters(tau=1)
 
-    def choose_action(self, observation, baseline : T.Tensor=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=}")
+        # 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 self.ig is not None:
+             attribution = self.ig.attribute(observation, baselines=baseline, n_steps=1)
+             print('Attributions:', attribution)
 
 
         mu_prime = mu + T.tensor(self.noise(), dtype=T.float).to(self.actor.device)

train.py

@@ -9,12 +9,15 @@ from captum.attr import (IntegratedGradients)
 class TrainingLoop:
     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
         }
 
+        defaults.update(**kwargs)
+
         self.env = gym.make(
             env_spec,
             **defaults
@@ -87,6 +90,7 @@ class TrainingLoop:
 
     def _collect_running_baseline_average(self, num_iterations: int) -> torch.Tensor:
         assert self.agent is not None
+        print("--------- Collecting running baseline average ----------")
 
         self.agent.load_models()
 
@@ -98,16 +102,19 @@ class TrainingLoop:
             obs, _ = self.env.reset()
             
             sum_obs += obs
-            print(f"Baseline on interation #{i}: {obs}")
+            # print(f"Baseline on interation #{i}: {obs}")
 
             while not done:
-                act = self.agent.choose_action(obs, attribution=None, baseline=None)
+                act = self.agent.choose_action(obs, baseline=None)
                 new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
                 score += reward
                 obs = new_state
 
+        print(f"Baseline collected: {sum_obs / num_iterations}")
+
         self.env.close()
+        
 
         return sum_obs / num_iterations
 
@@ -122,8 +129,12 @@ class TrainingLoop:
 
         baseline = baseline_options[option]
 
+        print("\n\n\n\n--------- Performing Attributions -----------")
+
         self.agent.load_models()
 
+        
+        print(self.agent.actor)
         ig = IntegratedGradients(self.agent.actor)
         self.agent.ig = ig
 
@@ -134,7 +145,7 @@ class TrainingLoop:
             score = 0
             obs, _ = self.env.reset()
             while not done:
-                act = self.agent.choose_action(obs, baseline=baseline)
+                act = self.agent.choose_action(observation=obs, baseline=baseline)
                 new_state, reward, terminated, truncated, info = self.env.step(act)
                 done = terminated or truncated
                 score += reward