Update K-Bot

sim/model_export.py

@@ -205,6 +205,52 @@ def forward(
 
         return actions_scaled, actions, x
 
+def get_actor_policy(model_path: str, cfg: ActorCfg) -> Tuple[nn.Module, dict, Tuple[Tensor, ...]]:
+    all_weights = torch.load(model_path, map_location="cpu", weights_only=True)
+    weights = all_weights["model_state_dict"]
+    num_actor_obs = weights["actor.0.weight"].shape[1]
+    num_critic_obs = weights["critic.0.weight"].shape[1]
+    num_actions = weights["std"].shape[0]
+    actor_hidden_dims = [v.shape[0] for k, v in weights.items() if re.match(r"actor\.\d+\.weight", k)]
+    critic_hidden_dims = [v.shape[0] for k, v in weights.items() if re.match(r"critic\.\d+\.weight", k)]
+    actor_hidden_dims = actor_hidden_dims[:-1]
+    critic_hidden_dims = critic_hidden_dims[:-1]
+
+    ac_model = ActorCritic(num_actor_obs, num_critic_obs, num_actions, actor_hidden_dims, critic_hidden_dims)
+    ac_model.load_state_dict(weights)
+
+    a_model = Actor(ac_model.actor, cfg)
+
+    # Gets the model input tensors.
+    x_vel = torch.randn(1)
+    y_vel = torch.randn(1)
+    rot = torch.randn(1)
+    t = torch.randn(1)
+    dof_pos = torch.randn(a_model.num_actions)
+    dof_vel = torch.randn(a_model.num_actions)
+    prev_actions = torch.randn(a_model.num_actions)
+    imu_ang_vel = torch.randn(3)
+    imu_euler_xyz = torch.randn(3)
+    buffer = a_model.get_init_buffer()
+    input_tensors = (x_vel, y_vel, rot, t, dof_pos, dof_vel, prev_actions, imu_ang_vel, imu_euler_xyz, buffer)
+
+    jit_model = torch.jit.script(a_model)
+
+    # Add sim2sim metadata
+    robot_effort = list(a_model.robot.effort().values())
+    robot_stiffness = list(a_model.robot.stiffness().values())
+    robot_damping = list(a_model.robot.damping().values())
+    num_actions = a_model.num_actions
+    num_observations = a_model.num_observations
+
+    return a_model, {
+        "robot_effort": robot_effort,
+        "robot_stiffness": robot_stiffness,
+        "robot_damping": robot_damping,
+        "num_actions": num_actions,
+        "num_observations": num_observations,
+    }, input_tensors
+
 
 def convert_model_to_onnx(model_path: str, cfg: ActorCfg, save_path: Optional[str] = None) -> ort.InferenceSession:
     """Converts a PyTorch model to a ONNX format.

sim/play.py

@@ -22,9 +22,10 @@
 
 from sim.env import run_dir  # noqa: E402
 from sim.envs import task_registry  # noqa: E402
-from sim.model_export import ActorCfg, convert_model_to_onnx  # noqa: E402
+from sim.model_export import ActorCfg, get_actor_policy  # noqa: E402
 from sim.utils.helpers import get_args  # noqa: E402
 from sim.utils.logger import Logger  # noqa: E402
+from kinfer.export.pytorch import export_to_onnx
 
 import torch  # isort: skip
 
@@ -81,8 +82,19 @@ def play(args: argparse.Namespace) -> None:
     # export policy as a onnx module (used to run it on web)
     if args.export_onnx:
         path = ppo_runner.alg.actor_critic
-        convert_model_to_onnx(path, ActorCfg(), save_path="policy.onnx")
-        print("Exported policy as onnx to: ", path)
+        policy_cfg = ActorCfg()
+        actor_model, sim2sim_info, input_tensors = get_actor_policy(path, policy_cfg)
+
+        # Merge policy_cfg and sim2sim_info into a single config object
+        export_config = {**vars(policy_cfg), **sim2sim_info}
+
+        policy = export_to_onnx(
+            actor_model,
+            input_tensors=input_tensors,
+            config=export_config,
+            save_path="kinfer_policy.onnx"
+        )
+        print("Exported policy as kinfer-compatible onnx to: ", path)
 
     # Prepare for logging
     env_logger = Logger(env.dt)

sim/requirements.txt

@@ -11,3 +11,5 @@ wandb
 tensorboard==2.14.0
 onnxscript
 # onnxruntime
+
+kinfer

sim/resources/stompypro/joints.py

@@ -127,12 +127,12 @@ def default_standing(cls) -> Dict[str, float]:
             Robot.legs.left.hip_pitch: -0.23,
             Robot.legs.left.hip_yaw: 0.0,
             Robot.legs.left.hip_roll: 0.0,
-            Robot.legs.left.knee_pitch: 0.441,
-            Robot.legs.left.ankle_pitch: -0.258,
+            Robot.legs.left.knee_pitch: -0.441,
+            Robot.legs.left.ankle_pitch: 0.258,
             Robot.legs.right.hip_pitch: -0.23,
             Robot.legs.right.hip_yaw: 0.0,
             Robot.legs.right.hip_roll: 0.0,
-            Robot.legs.right.knee_pitch: 0.441,
+            Robot.legs.right.knee_pitch: -0.441,
             Robot.legs.right.ankle_pitch: -0.258,
         }