cp

Author: daveey

agent/src/metta/agent/policies/vit.py

@@ -143,7 +143,12 @@ def _prediction_step(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Te
 
 
 class ViTDefaultConfig(PolicyArchitecture):
-    """Speed-optimized ViT variant with lighter token embeddings and attention stack."""
+    """Speed-optimized ViT variant with lighter token embeddings and attention stack.
+
+    The trunk uses Axon blocks (post-up experts with residual connections) for efficient
+    feature processing. Configure trunk depth, layer normalization, and hidden dimension
+    scaling independently.
+    """
 
     class_path: str = "metta.agent.policy_auto_builder.PolicyAutoBuilder"
 
@@ -156,6 +161,12 @@ class ViTDefaultConfig(PolicyArchitecture):
     pass_state_during_training: bool = False
     _critic_hidden = 512
 
+    # Trunk configuration
+    # Number of Axon layers in the trunk (default: 16 for large model)
+    trunk_num_resnet_layers: int = 1
+    # Enable layer normalization after each trunk layer
+    trunk_use_layer_norm: bool = True
+
     components: List[ComponentConfig] = [
         ObsShimTokensConfig(in_key="env_obs", out_key="obs_shim_tokens", max_tokens=48),
         ObsAttrEmbedFourierConfig(
@@ -181,9 +192,9 @@ class ViTDefaultConfig(PolicyArchitecture):
             key_prefix="vit_cortex_state",
             stack_cfg=build_cortex_auto_config(
                 d_hidden=_latent_dim,
-                num_layers=1,
-                pattern="L",
-                post_norm=False,
+                num_layers=trunk_num_resnet_layers,
+                pattern="A",  # Axon blocks provide residual-like connections
+                post_norm=trunk_use_layer_norm,
             ),
             pass_state_during_training=pass_state_during_training,
         ),
@@ -209,32 +220,20 @@ class ViTDefaultConfig(PolicyArchitecture):
     action_probs_config: ActionProbsConfig = ActionProbsConfig(in_key="logits")
 
     def make_policy(self, policy_env_info: PolicyEnvInterface) -> Policy:
-        # Ensure downstream components match core dimension
-        # (self._latent_dim might have been overridden on the instance without updating all components)
-        cortex = next(c for c in self.components if isinstance(c, CortexTDConfig))
-        core_dim = cortex.out_features or cortex.d_hidden
-
-        actor_mlp = next(c for c in self.components if c.name == "actor_mlp")
-        assert isinstance(actor_mlp, MLPConfig)
-        if actor_mlp.in_features != core_dim:
-            actor_mlp.in_features = core_dim
-
-        critic = next(c for c in self.components if c.name == "critic")
-        assert isinstance(critic, MLPConfig)
-        if critic.in_features != core_dim:
-            critic.in_features = core_dim
+        # Note: trunk configuration (num_layers, layer_norm, scaling) is applied
+        # via the components list definition above, no runtime modification needed
 
         AgentClass = load_symbol(self.class_path)
         policy = AgentClass(policy_env_info, self)
         policy.num_actions = policy_env_info.action_space.n
 
         # Dimensions
-        latent_dim = core_dim
-        num_actions = policy.num_actions
+        latent_dim = int(self._latent_dim)
+        num_actions = int(policy.num_actions)
 
         # Dynamics Model: (Hidden + Action) -> (Hidden + Reward)
-        dyn_input_dim = latent_dim + num_actions
-        dyn_output_dim = latent_dim + 1
+        dyn_input_dim = int(latent_dim + num_actions)
+        dyn_output_dim = int(latent_dim + 1)
 
         # Simple MLP for dynamics
         dynamics_net = nn.Sequential(
@@ -246,7 +245,7 @@ def make_policy(self, policy_env_info: PolicyEnvInterface) -> Policy:
 
         # Returns/Reward Prediction Heads (for Muesli)
         # Input: Core + Logits
-        pred_input_dim = latent_dim + num_actions
+        pred_input_dim = int(latent_dim + num_actions)
 
         returns_module = nn.Linear(pred_input_dim, 1)
         reward_module = nn.Linear(pred_input_dim, 1)