Merge pull request #101 from UT-Austin-RPL/rope

rope, `action_from_buffer`,  sigma reparam init

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
 
 

amago/nets/traj_encoders.py

@@ -345,7 +345,8 @@ class TformerTrajEncoder(TrajEncoder):
         activation: Activation function. Defaults to "leaky_relu".
         norm: Normalization function. Defaults to "layer" (LayerNorm).
         pos_emb: Position embedding type. "fixed" (default) uses sinusoidal
-            embeddings, "learned" uses trainable embeddings per timestep.
+            embeddings, "learned" uses trainable embeddings per timestep,
+            "rope" uses Rotary Position Embeddings applied to Q and K.
         causal: Whether to use causal attention mask. Defaults to True.
         sigma_reparam: Whether to use :math:`\sigma`-reparam feed-forward layers
             from https://arxiv.org/abs/2303.06296. Defaults to True.
@@ -395,6 +396,7 @@ def make_layer():
                     dropout_qkv=dropout_qkv,
                     head_scaling=head_scaling,
                     sigma_reparam=sigma_reparam,
+                    use_rope=pos_emb == "rope",
                 ),
                 d_model=self.d_model,
                 d_ff=d_ff,

amago/nets/transformer.py

@@ -423,8 +423,9 @@ def __init__(self, d_in, d_out, bias: bool = True, fast_init: bool = False):
         super().__init__(d_in, d_out, bias=bias)
         if not fast_init:
             nn.init.trunc_normal_(self.weight, std=0.02)
-            u = torch.linalg.svd(self.weight.T, full_matrices=False)[-1][0].detach()
-            v = torch.linalg.svd(self.weight, full_matrices=False)[-1][0].detach()
+            U, _S, Vh = torch.linalg.svd(self.weight, full_matrices=False)
+            u = U[:, 0].detach()
+            v = Vh[0].detach()
         else:
             # initialization from legacy version used in the original AMAGO paper.
             # This was a guess based on the sigma reparam pseudocode before the code was released,
@@ -469,10 +470,12 @@ def __init__(
         dropout_qkv: float = 0.0,
         head_scaling: bool = True,
         sigma_reparam: bool = True,
+        use_rope: bool = False,
     ):
         super().__init__()
         assert isinstance(self_attention, SelfAttention)
         self.self_attention = self_attention
+        self.rope = RotaryPosEmb(d_qkv) if use_rope else None
         FF = SigmaReparam if sigma_reparam else nn.Linear
         self.qkv_projection = FF(d_model, 3 * d_qkv * n_heads, bias=False)
         self.dropout_qkv = nn.Dropout(dropout_qkv)
@@ -482,14 +485,23 @@ def __init__(
         )
         self.n_heads = n_heads
 
-    def forward(self, sequence, key_cache=None, val_cache=None, cache_seqlens=None):
+    def forward(
+        self,
+        sequence,
+        key_cache=None,
+        val_cache=None,
+        cache_seqlens=None,
+        pos_idxs=None,
+    ):
         qkv = self.dropout_qkv(self.qkv_projection(sequence))
         qkv = rearrange(
             qkv,
             "batch len (three d_qkv heads) -> batch len three heads d_qkv",
             heads=self.n_heads,
             three=3,
         )
+        if self.rope is not None:
+            qkv = self.rope(qkv, pos_idxs)
         out = self.head_scaler * self.self_attention(
             qkv=qkv,
             key_cache=key_cache,
@@ -538,10 +550,21 @@ def __init__(
         self.d_model = d_model
 
     @torch.compile
-    def forward(self, self_seq, key_cache=None, val_cache=None, cache_seqlens=None):
+    def forward(
+        self,
+        self_seq,
+        key_cache=None,
+        val_cache=None,
+        cache_seqlens=None,
+        pos_idxs=None,
+    ):
         q1 = self.norm1(self_seq)  # pre-norm
         q1 = self.attention_layer(
-            q1, key_cache=key_cache, val_cache=val_cache, cache_seqlens=cache_seqlens
+            q1,
+            key_cache=key_cache,
+            val_cache=val_cache,
+            cache_seqlens=cache_seqlens,
+            pos_idxs=pos_idxs,
         )
         q1 = self.norm2(q1)  # normformer extra norm 1
         self_seq = self_seq + q1
@@ -667,6 +690,39 @@ def forward(self, pos_idxs: torch.LongTensor):
         return self.embeddings(pos_idxs)
 
 
+class RotaryPosEmb(nn.Module):
+    """Rotary Position Embedding (RoPE). Half-split (Llama-style) convention."""
+
+    def __init__(self, head_dim: int, base: float = 10000.0):
+        super().__init__()
+        assert head_dim % 2 == 0, f"RoPE requires even head_dim, got {head_dim}"
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.head_dim = head_dim
+
+    def forward(self, qkv: torch.Tensor, pos_idxs: torch.Tensor) -> torch.Tensor:
+        """Rotate Q and K in packed (B, L, 3, H, D) QKV; V is unchanged."""
+        pos = pos_idxs.squeeze(-1).to(self.inv_freq.dtype)
+        freqs = torch.einsum("bl,d->bld", pos, self.inv_freq)
+        cos = freqs.cos().unsqueeze(2)
+        sin = freqs.sin().unsqueeze(2)
+
+        q, k, v = qkv.unbind(dim=2)
+        q = self._apply_rotary(q, cos, sin).to(q.dtype)
+        k = self._apply_rotary(k, cos, sin).to(k.dtype)
+        return torch.stack([q, k, v], dim=2)
+
+    @staticmethod
+    def _apply_rotary(
+        x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        """x: (B, L, H, D), cos/sin: (B, L, 1, D/2)."""
+        d_half = x.shape[-1] // 2
+        x1 = x[..., :d_half]
+        x2 = x[..., d_half:]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+
 class Transformer(nn.Module):
     """Build a full Transformer model from a list of layers."""
 
@@ -680,13 +736,16 @@ def __init__(
         pos_emb: str = "fixed",
     ):
         super().__init__()
+        self.use_rope = pos_emb == "rope"
         if pos_emb == "fixed":
             self.position_embedding = FixedPosEmb(d_model)
         elif pos_emb == "learnable":
             self.position_embedding = LearnablePosEmb(d_model)
+        elif pos_emb == "rope":
+            self.position_embedding = None
         else:
             raise ValueError(
-                f"Unrecognized pos_emb: {pos_emb}. Options are 'fixed' or 'learnable'."
+                f"Unrecognized pos_emb: {pos_emb}. Options are 'fixed', 'learnable', or 'rope'."
             )
         self.inp = nn.Linear(inp_dim, d_model)
         self.dropout = nn.Dropout(dropout_emb)
@@ -701,20 +760,22 @@ def emb_dim(self):
         return self.d_model
 
     def preprocess_seq(self, seq, pos_idxs):
-        pos_emb = self.position_embedding(pos_idxs.squeeze(-1))
         traj_emb = self.inp(seq)
-        traj_emb = self.dropout(traj_emb + pos_emb)
+        if self.position_embedding is not None:
+            pos_emb = self.position_embedding(pos_idxs.squeeze(-1))
+            traj_emb = traj_emb + pos_emb
+        traj_emb = self.dropout(traj_emb)
         return traj_emb
 
     @torch.compile
-    def training_forward(self, seq):
+    def training_forward(self, seq, pos_idxs=None):
         for layer in self.layers:
-            seq = layer(seq)
+            seq = layer(seq, pos_idxs=pos_idxs)
         return self.norm(seq)
 
-    def inference_forward(self, seq, hidden_state):
+    def inference_forward(self, seq, hidden_state, pos_idxs=None):
         for i, layer in enumerate(self.layers):
-            seq = layer(seq, *hidden_state[i])
+            seq = layer(seq, *hidden_state[i], pos_idxs=pos_idxs)
         return self.norm(seq)
 
     def forward(self, seq, pos_idxs, hidden_state: Optional[TformerHiddenState] = None):
@@ -731,10 +792,11 @@ def forward(self, seq, pos_idxs, hidden_state: Optional[TformerHiddenState] = No
         """
 
         traj_emb = self.preprocess_seq(seq, pos_idxs)
+        rope_pos = pos_idxs if self.use_rope else None
         if hidden_state is not None:
             assert not self.training
-            traj_emb = self.inference_forward(traj_emb, hidden_state)
+            traj_emb = self.inference_forward(traj_emb, hidden_state, pos_idxs=rope_pos)
             hidden_state.update()
         else:
-            traj_emb = self.training_forward(traj_emb)
+            traj_emb = self.training_forward(traj_emb, pos_idxs=rope_pos)
         return traj_emb, hidden_state

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "amago"
-version = "3.2.0"
+version = "3.3.0"
 description = "off-policy RL on long sequences"
 readme = { file = "README.md", content-type = "text/markdown" }
 authors = [