Init

legged_gym/utils/exporter.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import copy
+import os
+import torch
+from torch import nn
+from typing import Optional
+
+
+def export_policy_as_jit(policy: object, path: str, normalizer: Optional[object] = None, filename="policy.pt"):
+    """Export policy into a Torch JIT file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported JIT file. Defaults to "policy.pt".
+    """
+    policy_exporter = _TorchPolicyExporter(policy, normalizer)
+    policy_exporter.export(path, filename)
+
+
+def export_policy_as_onnx(
+    policy: object, path: str, normalizer: Optional[object] = None, filename="policy.onnx", verbose=False
+):
+    """Export policy into a Torch ONNX file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported ONNX file. Defaults to "policy.onnx".
+        verbose: Whether to print the model summary. Defaults to False.
+    """
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+    policy_exporter = _OnnxPolicyExporter(policy, normalizer, verbose)
+    policy_exporter.export(path, filename)
+
+
+
+def export_policy_as_pkl(
+    policy: nn.Module, path: str, filename="policy.pkl"
+):
+    """Export policy into a Torch pkl file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported pkl file. Defaults to "policy.pkl".
+    """
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+    model_dict = policy.state_dict()
+    torch.save(model_dict, os.path.join(path, filename))
+
+
+"""
+Helper Classes - Private.
+"""
+
+
+class _TorchPolicyExporter(torch.nn.Module):
+    """Exporter of actor-critic into JIT file."""
+
+    def __init__(self, policy, normalizer=None):
+        super().__init__()
+        self.is_recurrent = policy.is_recurrent
+        # copy policy parameters
+        if hasattr(policy, "student_encoder"):
+            self.student_encoder = copy.deepcopy(policy.student_encoder).cpu()
+            self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
+            self.forward = self.forward_cts
+        if hasattr(policy, "student_moe_encoder"):
+            self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder).cpu()
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            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 hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+            if self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+        elif hasattr(policy, "student"):
+            self.actor = copy.deepcopy(policy.student)
+            if self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_s.rnn)
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+        # set up recurrent network
+        if self.is_recurrent:
+            self.rnn.cpu()
+            self.register_buffer("hidden_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
+            self.register_buffer("cell_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
+            self.forward = self.forward_lstm
+            self.reset = self.reset_memory
+        # copy normalizer if exists
+        if normalizer:
+            self.normalizer = copy.deepcopy(normalizer)
+        else:
+            self.normalizer = torch.nn.Identity()
+
+    def forward_lstm(self, x):
+        x = self.normalizer(x)
+        x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
+        self.hidden_state[:] = h
+        self.cell_state[:] = c
+        x = x.squeeze(0)
+        return self.actor(x)
+
+    def forward(self, x):
+        return self.actor(self.normalizer(x))
+    
+    def forward_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        latent = self.student_encoder(self.history.flatten(1))
+        x = torch.cat([latent, x], dim=1)
+        return self.actor(x)
+    
+    def forward_moe_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        history_no_goal = self.history.reshape(1, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(1, -1)
+        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)
+
+    @torch.jit.export
+    def reset(self):
+        if hasattr(self, 'history'):
+            self.history = torch.zeros_like(self.history)
+
+    def reset_memory(self):
+        self.hidden_state[:] = 0.0
+        self.cell_state[:] = 0.0
+
+    def export(self, path, filename):
+        os.makedirs(path, exist_ok=True)
+        path = os.path.join(path, filename)
+        self.to("cpu")
+        traced_script_module = torch.jit.script(self)
+        traced_script_module.save(path)
+
+
+class _OnnxPolicyExporter(torch.nn.Module):
+    """Exporter of actor-critic into ONNX file."""
+
+    def __init__(self, policy, normalizer=None, verbose=False):
+        super().__init__()
+        self.verbose = verbose
+        self.input_dim = None
+        self.num_actions = 12
+        self.normalizer = torch.nn.Identity()
+        
+        # copy policy parameters
+        if hasattr(policy, 'student_encoder'):
+            self.student_encoder = copy.deepcopy(policy.student_encoder)
+            self.forward = self.forward_cts
+            self.input_dim = self.student_encoder[0].in_features
+            
+        elif hasattr(policy, "student_moe_encoder"):
+            self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            self.history_length = policy.history.shape[1]
+            self.forward = self.forward_moe_cts
+            self.input_dim = self.history_length * policy.history.shape[2]
+            
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+            if hasattr(self, 'is_recurrent') and self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+            if self.input_dim is None:
+                 self.input_dim = self.actor[0].in_features
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+
+    def forward_cts(self, x):  # x is stack observations by terms
+        x = self.normalizer(x)
+        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        obs_dim = sum(term_dims)
+        if x.shape[1] % obs_dim != 0:
+            raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
+            
+        frames = x.shape[1] // obs_dim
+        split_sizes = [dim * frames for dim in term_dims]
+        # [B, dim0*frames], [B, dim1*frames], ...
+        term_chunks = torch.split(x, split_sizes, dim=1)
+
+        # [ [B, frames, dim0], [B, frames, dim1], ... ]
+        frame_terms_reshaped = [
+            chunk.view(-1, frames, dim) 
+            for chunk, dim in zip(term_chunks, term_dims)
+        ]
+
+        history_by_frame = []
+        for i in range(frames):
+            # [ [B, dim0], [B, dim1], ... ]
+            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
+            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
+        # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
+        history = torch.cat(history_by_frame, dim=1)
+
+        last_obs = history[:, -obs_dim:]
+        latent = self.student_encoder(history)
+        x = torch.cat([latent, last_obs], dim=1)
+        
+        return self.actor(x)
+
+    def forward_moe_cts(self, x):
+        x = self.normalizer(x)
+        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        obs_dim = sum(term_dims)
+        if x.shape[1] % obs_dim != 0:
+            raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
+            
+        frames = x.shape[1] // obs_dim
+        split_sizes = [dim * frames for dim in term_dims]
+        # [B, dim0*frames], [B, dim1*frames], ...
+        term_chunks = torch.split(x, split_sizes, dim=1)
+
+        # [ [B, frames, dim0], [B, frames, dim1], ... ]
+        frame_terms_reshaped = [
+            chunk.view(-1, frames, dim) 
+            for chunk, dim in zip(term_chunks, term_dims)
+        ]
+
+        history_by_frame = []
+        for i in range(frames):
+            # [ [B, dim0], [B, dim1], ... ]
+            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
+            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
+        # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
+        history = torch.cat(history_by_frame, dim=1)
+
+        last_obs = history[:, -obs_dim:]
+        history_3d = history.view(-1, self.history_length, obs_dim)
+        history_no_goal = history_3d[:, :, self.obs_no_goal_mask].reshape(x.shape[0], -1)
+
+        latent, weights = self.student_moe_encoder(history, history_no_goal)
+        x = torch.cat([latent, last_obs], dim=1)
+
+        return self.actor(x), weights, latent
+
+    def export(self, path, filename):
+        self.to("cpu")
+        obs = torch.zeros(1, self.input_dim)
+        
+        output_names = ["actions"]
+        if self.forward == self.forward_moe_cts:
+            output_names.append("weights")
+            output_names.append("latent")
+
+        torch.onnx.export(
+            self,
+            obs,
+            os.path.join(path, filename),
+            export_params=True,
+            opset_version=11,
+            verbose=self.verbose,
+            input_names=["obs"],
+            output_names=output_names,
+            dynamic_axes={},
+        )