Initial commit

demo/export_skrl_policy_to_onnx.py

@@ -0,0 +1,185 @@
+"""
+Export SKRL PyTorch policy to ONNX format.
+
+Usage:
+    python demo/export_skrl_policy_to_onnx.py --checkpoint <path_to_checkpoint.pt> --output <output_dir>
+
+Example:
+    python demo/export_skrl_policy_to_onnx.py \
+        --checkpoint runs/go1-rough-terrain-walk/skrl/[time]_PPO/checkpoints/best_agent.pt \
+        --output exports_go1_rough
+"""
+
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# ============== PyTorch Checkpoint Export ==============
+
+class SKRLPolicyTorch(nn.Module):
+    """PyTorch policy matching SKRL GaussianMixin architecture.
+
+    The architecture is inferred from the checkpoint keys:
+    - net.N: Hidden layers (Linear -> ELU)
+    - mean_layer: Output layer (no activation)
+    """
+
+    def __init__(self, obs_dim, action_dim, hidden_dims=[512, 256, 128]):
+        super().__init__()
+        self.obs_dim = obs_dim
+        self.action_dim = action_dim
+        self.hidden_dims = hidden_dims
+
+        layers = []
+        in_dim = obs_dim
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(in_dim, hidden_dim),
+                nn.ELU(),
+            ])
+            in_dim = hidden_dim
+
+        self.net = nn.Sequential(*layers)
+        self.mean_layer = nn.Linear(in_dim, action_dim)
+        self.log_std = nn.Parameter(torch.zeros(action_dim))
+
+    def forward(self, x):
+        return self.mean_layer(self.net(x))
+
+
+class ONNXPolicyExporterTorch(nn.Module):
+    """Exporter wrapper that applies normalization and policy.
+
+    Matches SKRL's RunningStandardScaler behavior:
+    - Normalize: (x - mean) / sqrt(var + eps)
+    - Clip to [-5.0, 5.0] (clip_threshold)
+    """
+
+    def __init__(self, policy, normalizer_mean, normalizer_std):
+        super().__init__()
+        self.policy = policy
+        # Register normalizer buffers as float32 (matching RunningStandardScaler)
+        self.register_buffer('mean', normalizer_mean.float())
+        self.register_buffer('std', normalizer_std.float())
+        self.clip_threshold = 5.0
+
+    def forward(self, x):
+        # Normalize (RunningStandardScaler._compute)
+        x = (x - self.mean) / (self.std + 1e-8)
+        # Clip to RunningStandardScaler clip_threshold
+        x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
+        return self.policy(x)
+
+
+def infer_architecture_from_state_dict(state_dict):
+    """Infer the policy architecture from state dict keys.
+
+    Args:
+        state_dict: Policy state dict
+
+    Returns:
+        tuple: (obs_dim, action_dim, hidden_dims)
+    """
+    # Find observation dimension from first layer
+    obs_dim = state_dict['net.0.weight'].shape[1]
+
+    # Find action dimension from mean_layer
+    action_dim = state_dict['mean_layer.weight'].shape[0]
+
+    # Infer hidden dimensions from net layers
+    hidden_dims = []
+    net_keys = sorted([k for k in state_dict.keys() if k.startswith('net.') and k.endswith('.weight')])
+    for key in net_keys:
+        if 'mean_layer' not in key:  # Skip output layer
+            layer_idx = int(key.split('.')[1])
+            if layer_idx % 2 == 0:  # Only Linear layers, not activation layers
+                hidden_dims.append(state_dict[key].shape[0])
+
+    return obs_dim, action_dim, hidden_dims
+
+
+def export_torch_policy_to_onnx(checkpoint_path, output_path, verbose=False):
+    """Export PyTorch SKRL policy to ONNX.
+
+    Args:
+        checkpoint_path: Path to .pt checkpoint file
+        output_path: Directory to save ONNX file
+        verbose: Whether to print verbose output
+    """
+    print(f"Loading PyTorch checkpoint from: {checkpoint_path}")
+    state = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+
+    policy_state = state['policy']
+    normalizer_state = state['state_preprocessor']
+
+    # Infer architecture
+    obs_dim, action_dim, hidden_dims = infer_architecture_from_state_dict(policy_state)
+    print(f"Inferred architecture: obs={obs_dim}, action={action_dim}, hidden={hidden_dims}")
+
+    # Create model
+    print("Creating PyTorch model...")
+    policy = SKRLPolicyTorch(obs_dim, action_dim, hidden_dims)
+
+    # Remove value layers from state dict (keep only policy)
+    policy_state_policy_only = {k: v for k, v in policy_state.items()
+                                if not k.startswith('value') and 'value_layer' not in k}
+    policy.load_state_dict(policy_state_policy_only, strict=False)
+    policy.eval()
+
+    # Load normalizer
+    mean = normalizer_state['running_mean']
+    std = torch.sqrt(normalizer_state['running_variance'])
+
+    # Create exporter
+    exporter = ONNXPolicyExporterTorch(policy, mean, std)
+    exporter.eval()
+
+    # Create output directory
+    os.makedirs(output_path, exist_ok=True)
+    output_file = os.path.join(output_path, "policy.onnx")
+
+    # Export to ONNX
+    print(f"Exporting policy to: {output_file}")
+    obs = torch.zeros(1, obs_dim, dtype=torch.float32)
+
+    torch.onnx.export(
+        exporter,
+        obs,
+        output_file,
+        export_params=True,
+        opset_version=11,
+        verbose=verbose,
+        input_names=["observations"],
+        output_names=["actions"],
+        dynamic_axes={},
+    )
+
+    print(f"Successfully exported policy to ONNX: {output_file}")
+
+    # Save PyTorch model and normalizer
+    torch.save(policy.state_dict(), os.path.join(output_path, "policy.pt"))
+    np.savez(os.path.join(output_path, "normalizer.npz"),
+             mean=mean.numpy(), std=std.numpy())
+
+    return output_file
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Export SKRL PyTorch policy to ONNX")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint file (.pt)")
+    parser.add_argument("--output", type=str, default="exports", help="Output directory")
+    parser.add_argument("--verbose", action="store_true", help="Verbose output")
+
+    args = parser.parse_args()
+
+    output_file = export_torch_policy_to_onnx(
+        checkpoint_path=args.checkpoint,
+        output_path=args.output,
+        verbose=args.verbose
+    )
+
+    print(f"\nExported to: {output_file}")