a_buffer critic mismatch at clipped extremes

Author: jakegrigsby

amago/agent.py

@@ -795,6 +795,7 @@ def get_values(self, batch: Batch) -> dict[str, torch.Tensor]:
         K_a = self.num_actions_for_value_in_actor_loss if not self.discrete else 1
         a_buffer = F.pad(a, (0, 0, 0, 1), "replicate")
         a_buffer = repeat(a_buffer, f"b l a -> b l {G} a")
+        a_buffer = self.actor.policy_dist.action_from_buffer(a_buffer)
         state_mask = (~((batch.rl2s == self.pad_val).all(-1, keepdim=True))).bool()[
             :, 1:, ...
         ]
@@ -927,6 +928,7 @@ def forward(self, batch: Batch, log_step: bool) -> torch.Tensor:
         # we give it a fake one to make shape math work.
         a_buffer = F.pad(a, (0, 0, 0, 1), "replicate")
         a_buffer = repeat(a_buffer, f"b l a -> b l {G} a")
+        a_buffer_critic = self.actor.policy_dist.action_from_buffer(a_buffer)
         C = len(self.critics)
         # arrays used by critic update end up in a (B, L, C, G, dim) format
         assert batch.rews.shape == (B, L - 1, 1)
@@ -964,7 +966,7 @@ def forward(self, batch: Batch, log_step: bool) -> torch.Tensor:
             ############################
             s_a_agent_g = (s_rep.detach(), a_agent)
             q_s_a_agent_g = self.maximized_critics(*s_a_agent_g).mean(0)
-            s_a_g = (s_rep[:, :-1, ...], a_buffer[:, :-1, ...].unsqueeze(0))
+            s_a_g = (s_rep[:, :-1, ...], a_buffer_critic[:, :-1, ...].unsqueeze(0))
             q_s_a_g = self.critics(*s_a_g, log_dict=active_log_dict).mean(0)
             assert q_s_a_agent_g.shape == (B, L, C, G, 1)
             assert q_s_a_g.shape == (B, L-1, C, G, 1)
@@ -1078,7 +1080,6 @@ def masked_avg(x_, dim=0):
             stats[f"Q(s, a) (global mean, rescaled) gamma={gamma:.3f}"] = masked_avg(
                 q_s_a_g, i
             )
-            print("here")
             stats["Q Sequence"] = q_s_a_g
             stats[f"Q(s,a) (global mean, raw scale) gamma={gamma:.3f}"] = masked_avg(
                 raw_q_s_a_g, i
@@ -1299,6 +1300,7 @@ def get_values(self, batch: Batch) -> dict[str, torch.Tensor]:
         K_a = self.num_actions_for_value_in_actor_loss
         a_buffer = F.pad(a, (0, 0, 0, 1), "replicate")
         a_buffer = repeat(a_buffer, f"b l a -> b l {G} a")
+        a_buffer = self.actor.policy_dist.action_from_buffer(a_buffer)
         state_mask = (~((batch.rl2s == self.pad_val).all(-1, keepdim=True))).bool()[
             :, 1:, ...
         ]
@@ -1363,6 +1365,7 @@ def forward(self, batch: Batch, log_step: bool):
         K_c = self.num_actions_for_value_in_critic_loss
         a_buffer = F.pad(a, (0, 0, 0, 1), "replicate")
         a_buffer = repeat(a_buffer, f"b l a -> b l {G} a")
+        a_buffer = self.actor.policy_dist.action_from_buffer(a_buffer)
         C = len(self.critics)
         assert batch.rews.shape == (B, L - 1, 1)
         assert batch.dones.shape == (B, L - 1, 1)

amago/nets/policy_dists.py

@@ -211,6 +211,15 @@ def input_dimension(self) -> int:
         """
         raise NotImplementedError
 
+    def action_from_buffer(self, action: torch.Tensor) -> torch.Tensor:
+        """Preprocess raw actions loaded from the replay buffer for use in the critic.
+
+        Subclasses override this to align buffer actions with the form the policy
+        actually produces (e.g., softening hard one-hot discrete actions).
+        Default is identity.
+        """
+        return action
+
     @abstractmethod
     def forward(
         self, vec: torch.Tensor, log_dict: Optional[dict] = None
@@ -269,6 +278,13 @@ def is_discrete(self) -> bool:
     def input_dimension(self) -> int:
         return self.d_action
 
+    def clip_and_normalize(self, probs: torch.Tensor) -> torch.Tensor:
+        clipped = probs.clamp(self.clip_prob_low, self.clip_prob_high)
+        return clipped / clipped.sum(-1, keepdim=True)
+
+    def action_from_buffer(self, action: torch.Tensor) -> torch.Tensor:
+        return self.clip_and_normalize(action)
+
     def forward(
         self, vec: torch.Tensor, log_dict: Optional[dict] = None
     ) -> _Categorical:
@@ -278,7 +294,7 @@ def forward(
         safe_probs = clip_probs / clip_probs.sum(-1, keepdims=True).detach()
         safe_dist = _Categorical(probs=safe_probs)
         if log_dict is not None:
-            add_activation_log("Discrete-probs", probs, log_dict)
+            add_activation_log("Discrete-probs", dist.probs, log_dict)
         return safe_dist