Add robogauge eval; update single obs export jit

scripts/rsl_rl/utils.py

@@ -6,7 +6,6 @@ import torch
 import re
 import os
 import sys
-from typing import NamedTuple
 
 "Script to log terminal output to a file, stripping ANSI escape codes."
 class Logger:
@@ -28,11 +27,6 @@ class Logger:
         self.terminal.flush()
         self.log.flush()
 
-# Inputs of CTS Policy is a TensorDict with 'policy' and 'single_obs' keys, we simulate this with a NamedTuple.
-class CTSPolicyInputs(NamedTuple):
-    policy: torch.Tensor
-    single_obs: torch.Tensor
-
 def export_cts_policy_as_jit(policy: object, actor_obs_normalizer: object | None, single_obs_normalizer: object | None, path: str, filename="policy.pt"):
     """Export CTS policy into a Torch JIT file.
 
@@ -75,6 +69,16 @@ class _TorchPolicyExporter(torch.nn.Module):
     """Exporter of actor-critic into JIT file."""
 
     def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None):
+        """Initialize a TorchScript exporter for CTS policy inference.
+
+        The exported model consumes only the current `single_obs` frame and maintains
+        internal stacked history to reconstruct actor observations expected by the student encoder.
+
+        Args:
+            policy: Source CTS policy module to export.
+            actor_obs_normalizer: Normalizer applied to stacked actor observations.
+            single_obs_normalizer: Normalizer applied to current single-frame observations.
+        """
         assert not policy.is_recurrent, "CTS policy should not be recurrent"
         super().__init__()
         
@@ -87,6 +91,22 @@ class _TorchPolicyExporter(torch.nn.Module):
             raise ValueError("Policy does not have an actor/student module.")
         self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
         self.state_dependent_std = policy.state_dependent_std
+        self.num_actions = int(policy.num_actions)
+        self.num_single_obs = int(policy.num_single_obs)
+        self.num_actor_obs = int(policy.num_actor_obs)
+        if self.num_actor_obs % self.num_single_obs != 0:
+            raise ValueError(
+                f"num_actor_obs ({self.num_actor_obs}) must be divisible by num_single_obs ({self.num_single_obs})."
+            )
+        self.history_len = self.num_actor_obs // self.num_single_obs
+        # Keep the same per-term history layout as deploy-side push_obs_history:
+        # [ang_vel(3), gravity(3), cmd(3), joint_pos(A), joint_vel(A), last_action(A)].
+        self.feature_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        if sum(self.feature_dims) != self.num_single_obs:
+            raise ValueError(
+                "Unsupported single_obs layout: expected 3+3+3+3*num_actions to match num_single_obs."
+            )
+        self.register_buffer("obs_history", torch.zeros(1, self.num_actor_obs, dtype=torch.float32))
         
         # copy normalizer if exists
         if actor_obs_normalizer:
@@ -98,9 +118,43 @@ class _TorchPolicyExporter(torch.nn.Module):
         else:
             self.single_obs_normalizer = torch.nn.Identity()
 
-    def forward(self, x: CTSPolicyInputs):
-        single_obs = self.single_obs_normalizer(x.single_obs)
-        obs_a = self.actor_obs_normalizer(x.policy)
+    def forward(self, single_obs: torch.Tensor):
+        """Compute policy action from one current observation frame.
+
+        The exporter keeps an internal FIFO history buffer and shifts it by one frame
+        on each forward call before appending the latest observation.
+
+        Args:
+            single_obs: Current-step observation tensor with shape `[B, num_single_obs]`.
+
+        Returns:
+            The policy action tensor.
+        """
+        if single_obs.dim() == 1:
+            single_obs = single_obs.unsqueeze(0)
+        if single_obs.shape[-1] != self.num_single_obs:
+            raise ValueError(
+                f"Expected single_obs last dimension {self.num_single_obs}, got {single_obs.shape[-1]}."
+            )
+        if single_obs.shape[0] != 1:
+            raise ValueError("TorchScript CTS deployment currently supports batch size 1 only.")
+
+        next_history = self.obs_history.clone()
+        history_offset = 0
+        single_offset = 0
+        for dim in self.feature_dims:
+            block_size = dim * self.history_len
+            block_end = history_offset + block_size
+            single_end = single_offset + dim
+            block = self.obs_history[:, history_offset:block_end]
+            shifted_block = torch.cat([block[:, dim:], single_obs[:, single_offset:single_end]], dim=-1)
+            next_history[:, history_offset:block_end] = shifted_block
+            history_offset = block_end
+            single_offset = single_end
+        self.obs_history.copy_(next_history)
+
+        single_obs = self.single_obs_normalizer(single_obs)
+        obs_a = self.actor_obs_normalizer(self.obs_history)
         latent, _ = self.student_moe_encoder(obs_a)
         latent_and_obs = torch.cat([latent, single_obs], dim=-1)
         if self.state_dependent_std:
@@ -110,7 +164,8 @@ class _TorchPolicyExporter(torch.nn.Module):
         
     @torch.jit.export
     def reset(self):
-        pass
+        """Reset internal observation history state."""
+        self.obs_history.zero_()
 
     def export(self, path, filename):
         os.makedirs(path, exist_ok=True)