Merge pull request #102 from UT-Austin-RPL/zs

minor changes to enable new methods

Author: jakegrigsby

amago/nets/ff.py

@@ -17,9 +17,11 @@ class Normalization(nn.Module):
 
     Args:
         method: Normalization method to use. Options are: "layer", "batch",
-            "rmsnorm", "unitball", "unitball-detach", "none". "unitball" is
-            (x / ||x||), "unitball-detach" is (x / ||x||.detach()). "none" is a
-            no-op and the rest are standard LayerNorm, BatchNorm, RMSNorm.
+            "rmsnorm", "unitball", "unitball-detach", "hypersphere", "none".
+            "unitball" is (x / ||x||), "unitball-detach" is (x / ||x||.detach()),
+            "hypersphere" is (x / ||x|| * sqrt(d)) projecting onto S^{d-1} with
+            ||x|| = sqrt(d). "none" is a no-op and the rest are standard
+            LayerNorm, BatchNorm, RMSNorm.
         d_model: Expected dimension of the input to normalize (scalar). Operates
             on the last dimensions of the input sequence.
     """
@@ -43,6 +45,8 @@ def __init__(self, method: Optional[str], d_model: int):
                     torch.linalg.vector_norm(x, ord=2, dim=-1, keepdim=True) + 1e-5
                 ).detach()
             )
+        elif method == "hypersphere":
+            self.norm = lambda x: F.normalize(x, dim=-1) * (x.shape[-1] ** 0.5)
         elif method == "rmsnorm":
             self.norm = _RMSNorm(size=d_model)
         elif method == "simnorm":

amago/nets/traj_encoders.py

@@ -164,6 +164,26 @@ def forward(
         pass
 
 
+@gin.configurable
+@register_traj_encoder("identity")
+class IdentityTrajEncoder(TrajEncoder):
+    """Passthrough trajectory encoder that returns tstep embeddings unchanged."""
+
+    @property
+    def emb_dim(self) -> int:
+        return self.tstep_dim
+
+    def forward(
+        self,
+        seq: torch.Tensor,
+        time_idxs: torch.Tensor,
+        hidden_state: Optional[Any] = None,
+        log_dict: Optional[dict] = None,
+    ) -> Tuple[torch.Tensor, Optional[Any]]:
+        new_hidden = None if hidden_state is None else hidden_state
+        return seq, new_hidden
+
+
 @gin.configurable
 @register_traj_encoder("ff")
 class FFTrajEncoder(TrajEncoder):

amago/nets/tstep_encoders.py

@@ -175,6 +175,9 @@ class FFTstepEncoder(TstepEncoder):
             space. If None, every key in the observation is used. Multi-modal
             observations are handled by flattening and concatenating values in a
             consistent order (alphabetical by key). Defaults to None.
+        auto_scale: If True, d_hidden and d_output are clamped to be at least
+            as large as the flattened input dimension, preventing bottlenecks
+            when observation keys contribute many features. Defaults to False.
     """
 
     def __init__(
@@ -190,6 +193,7 @@ def __init__(
         hide_rl2s: bool = False,
         normalize_inputs: bool = True,
         specify_obs_keys: Optional[list[str]] = None,
+        auto_scale: bool = False,
     ):
         super().__init__(obs_space=obs_space, rl2_space=rl2_space, hide_rl2s=hide_rl2s)
         if specify_obs_keys is None:
@@ -200,6 +204,9 @@ def __init__(
             math.prod(self.obs_space[key].shape) for key in self.obs_keys
         )
         in_dim = flat_obs_shape + self.rl2_space.shape[-1]
+        if auto_scale:
+            d_hidden = max(d_hidden, in_dim)
+            d_output = max(d_output, in_dim)
         self.in_norm = InputNorm(in_dim, skip=not normalize_inputs)
         self.base = ff.MLP(
             d_inp=in_dim,