v0.1.6; add rem-cts

legged_gym/utils/exporter.py

@@ -82,6 +82,8 @@ class _TorchPolicyExporter(torch.nn.Module):
             self.history_length = policy.history.shape[1]
             self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
             self.forward = self.forward_moe_cts
+            if not hasattr(policy, "obs_no_goal_mask"):
+                self.forward = self.forward_rem_cts
         if hasattr(policy, "actor_mcp"):
             self.actor = copy.deepcopy(policy.actor_mcp)
             self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
@@ -139,6 +141,13 @@ class _TorchPolicyExporter(torch.nn.Module):
         latent, weights = self.student_moe_encoder(self.history.flatten(1), history_no_goal)
         x = torch.cat([latent, x], dim=1)
         return self.actor(x), (weights, latent)
+
+    def forward_rem_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        latent, weights = self.student_moe_encoder(self.history.flatten(1))
+        x = torch.cat([latent, x], dim=1)
+        return self.actor(x), (weights, latent)
     
     def forward_mcp_cts(self, x):  # x is single observations
         x = self.normalizer(x)