Merge branch 'kinfer-schema' of https://github.com/kscalelabs/sim int…

…o kinfer-schema
kscalelabs · Jan 2, 2025 · dfd2263 · dfd2263
2 parents d3ca3d0 + 34783c0
commit dfd2263
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/sim/envs/base/legged_robot.py b/sim/envs/base/legged_robot.py
@@ -471,7 +471,6 @@ def update_command_curriculum(self, env_ids):
                 self.command_ranges["lin_vel_x"][1] + 0.5, 0.0, self.cfg.commands.max_curriculum
             )
 
-
     def _resample_default_positions(self):
         self.default_dof_pos = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False)
         for i in range(self.num_dofs):

diff --git a/sim/envs/humanoids/gpr_config.py b/sim/envs/humanoids/gpr_config.py
@@ -1,5 +1,7 @@
 """Defines the environment configuration for the Getting up task"""
 
+from kinfer import proto as P
+
 from sim.env import robot_urdf_path
 from sim.envs.base.legged_robot_config import (  # type: ignore
     LeggedRobotCfg,
@@ -26,8 +28,6 @@ class env(LeggedRobotCfg.env):
         episode_length_s = 24  # episode length in seconds
         use_ref_actions = False
 
-        from kinfer import proto as P
-
         input_schema = P.IOSchema(
             values=[
                 P.ValueSchema(
@@ -169,7 +169,7 @@ class noise_scales:
             height_measurements = 0.1
 
         class noise_ranges:
-            default_pos = 0.03 # +- 0.05 rad
+            default_pos = 0.03  # +- 0.05 rad
 
     class init_state(LeggedRobotCfg.init_state):
         pos = [0.0, 0.0, Robot.height]

diff --git a/sim/sim2sim.py b/sim/sim2sim.py
@@ -392,17 +392,23 @@ def run_mujoco(
     policy = ONNXModel(args.load_model)
     metadata = policy.attached_metadata
 
+    joint_names = []
+    for value_schema in policy.input_schema.values:
+        if value_schema.HasField("joint_positions"):
+            joint_names = list(value_schema.joint_positions.joint_names)
+            break
+
     try:
         model_info = {
-            "num_actions": metadata["num_actions"],
+            "num_actions": len(joint_names),
             "num_observations": metadata["num_observations"],
-            "robot_effort": [metadata["robot_effort"][joint] for joint in metadata["joint_names"]],
-            "robot_stiffness": [metadata["robot_stiffness"][joint] for joint in metadata["joint_names"]],
-            "robot_damping": [metadata["robot_damping"][joint] for joint in metadata["joint_names"]],
+            "robot_effort": [metadata["robot_effort"][joint] for joint in joint_names],
+            "robot_stiffness": [metadata["robot_stiffness"][joint] for joint in joint_names],
+            "robot_damping": [metadata["robot_damping"][joint] for joint in joint_names],
             "sim_dt": metadata["sim_dt"],
             "sim_decimation": metadata["sim_decimation"],
             "tau_factor": metadata["tau_factor"],
-            "joint_names": metadata["joint_names"],
+            "joint_names": joint_names,
             "cycle_time": metadata["cycle_time"],
         }
     except Exception as e: