Fix warp-lang & remove sensor re-naming

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
   - id: trailing-whitespace
   - id: debug-statements
     exclude: ^molmo_spaces/utils/misc_utils\.py$
-    language_version: python3.10
+    language_version: python3.11
   - id: pretty-format-json
     args: ["--autofix", "--indent=4", "--no-sort-keys"]
 

docs/sensors.md

@@ -117,7 +117,7 @@ class MyTask(BaseMujocoTask):
         sensors = get_core_sensors(config)
         sensors.extend([
             ObjectStartPoseSensor(object_name=config.task_config.pickup_obj_name,
-                                  uuid="obj_start_pose"),
+                                  uuid="obj_start"),
             GraspStateSensor(object_name=config.task_config.pickup_obj_name,
                              uuid="grasp_state_pickup_obj"),
         ])
@@ -228,7 +228,7 @@ These are added by individual tasks on top of the core suite. Notable examples:
 
 | Task | Adds (on top of `get_core_sensors`) |
 |---|---|
-| [`PickTask`][molmo_spaces.tasks.pick_task.PickTask] | `ObjectStartPoseSensor(uuid="obj_start_pose")`, `GraspStateSensor(uuid="grasp_state_pickup_obj")`, `PickupObjGoalPoseSensor(uuid="obj_end_pose")` |
+| [`PickTask`][molmo_spaces.tasks.pick_task.PickTask] | `ObjectStartPoseSensor(uuid="obj_start")`, `GraspStateSensor(uuid="grasp_state_pickup_obj")`, `PickupObjGoalPoseSensor(uuid="obj_end")` |
 | [`PickAndPlaceTask`][molmo_spaces.tasks.pick_and_place_task.PickAndPlaceTask] | `ObjectStartPoseSensor`, `GraspStateSensor` for both the pickup object **and** the place receptacle |
 | [`OpeningTask`][molmo_spaces.tasks.opening_tasks.OpeningTask] | Inherits the `PickTask` suite (since opening an articulated object is structurally similar to picking) |
 | [`DoorOpeningTask`][molmo_spaces.tasks.opening_tasks.DoorOpeningTask] | Uses `get_rby1_door_opening_sensors` — see warning below |

molmo_spaces/kinematics/parallel/warp_kinematics.py

@@ -2,17 +2,17 @@
 Provides a general-purpose, robot-agnostic, vectorized (and optionally GPU accelerated) kinematics solver.
 """
 
+import logging
+from collections import OrderedDict
 from dataclasses import dataclass
 from functools import cache
 from typing import TYPE_CHECKING
-from collections import OrderedDict
-import logging
 
-import numpy as np
 import mujoco
-from mujoco import MjSpec, MjModel, MjData
 import mujoco_warp as mjw
+import numpy as np
 import warp as wp
+from mujoco import MjData, MjModel, MjSpec
 
 from molmo_spaces.kinematics.parallel.parallel_kinematics import ParallelKinematics
 from molmo_spaces.robots.robot_views.abstract import (
@@ -166,7 +166,7 @@ def cholesky_solve6(H: mat66f, b: vec6f) -> vec6f:
     L = mat66f()
     for i in range(6):
         for j in range(i + 1):
-            s = float(0.0)
+            s = wp.float32(0.0)
             for k in range(j):
                 s += L[i, k] * L[j, k]
             if i == j:
@@ -177,15 +177,15 @@ def cholesky_solve6(H: mat66f, b: vec6f) -> vec6f:
     # Forward substitution: L @ y = b
     y = vec6f()
     for i in range(6):
-        s = float(0.0)
+        s = wp.float32(0.0)
         for k in range(i):
             s += L[i, k] * y[k]
         y[i] = (b[i] - s) / L[i, i]
 
     # Backward substitution: L^T @ x = y
     x = vec6f()
     for i in range(5, -1, -1):
-        s = float(0.0)
+        s = wp.float32(0.0)
         for k in range(i + 1, 6):
             s += L[k, i] * x[k]
         x[i] = (y[i] - s) / L[i, i]
@@ -223,10 +223,16 @@ def lm_step(
     H = mat66f()
     for a in range(6):
         for b in range(6):
-            val = float(0.0)
+            val = wp.float32(0.0)
             for k in range(nv):
-                Ja = jacp[i, a, k] if a < 3 else jacr[i, a - 3, k]
-                Jb = jacp[i, b, k] if b < 3 else jacr[i, b - 3, k]
+                if a < 3:
+                    Ja = jacp[i, a, k]
+                else:
+                    Ja = jacr[i, a - 3, k]
+                if b < 3:
+                    Jb = jacp[i, b, k]
+                else:
+                    Jb = jacr[i, b - 3, k]
                 val += Ja * Jb
             if a == b:
                 val += damping[0]
@@ -237,7 +243,7 @@ def lm_step(
 
     # q_dot = J^T @ x, dq = q_dot * dt
     for k in range(nv):
-        val = float(0.0)
+        val = wp.float32(0.0)
         for a in range(3):
             val += jacp[i, a, k] * x[a]
             val += jacr[i, a, k] * x[a + 3]
@@ -281,11 +287,7 @@ def __init__(self, robot_config: "BaseRobotConfig", device: str = "cpu"):
         self._frame_move_groups: dict[str, MJCFFrameMixin] = {}
         for mg_id in self._robot_view.move_group_ids():
             mg = self._robot_view.get_move_group(mg_id)
-            assert (
-                mg.n_joints == 0
-                or isinstance(mg, SimplyActuatedMoveGroup)
-                or isinstance(mg, GripperGroup)
-            )
+            assert mg.n_joints == 0 or isinstance(mg, (SimplyActuatedMoveGroup, GripperGroup))
             if isinstance(mg, SimplyActuatedMoveGroup):
                 self._actuated_move_groups[mg_id] = mg
             if isinstance(mg, MJCFFrameMixin):
@@ -386,7 +388,7 @@ def fk(
 
         qpos_arr = self._dicts_to_qpos_arr(qpos_dicts)
         with wp.ScopedDevice(self._device):
-            wp.copy(data.qpos, wp.from_numpy(qpos_arr))
+            wp.copy(data.qpos, wp.from_numpy(qpos_arr, dtype=wp.float32))
             mjw.fwd_position(self._mjw_model, data)
 
         dol = {}
@@ -605,11 +607,13 @@ def ik(
             ik_args.dt.fill_(dt)
             wp.copy(
                 ik_args.jacobian_mask,
-                wp.from_numpy(self._create_jacobian_mask(batch_size, unlocked_move_group_ids)),
+                wp.from_numpy(
+                    self._create_jacobian_mask(batch_size, unlocked_move_group_ids), dtype=wp.int32
+                ),
             )
 
             q0_arr = self._dicts_to_qpos_arr(q0_dicts)
-            wp.copy(data.qpos, wp.from_numpy(q0_arr))
+            wp.copy(data.qpos, wp.from_numpy(q0_arr, dtype=wp.float32))
 
             for i in range(max_iter):
                 if self._device.startswith("cuda"):

molmo_spaces/tasks/pick_and_place_task.py

@@ -51,7 +51,7 @@ def _create_sensor_suite_from_config(self, config: MlSpacesExpConfig) -> SensorS
         sensors.extend(
             [
                 ObjectStartPoseSensor(
-                    object_name=config.task_config.pickup_obj_name, uuid="obj_start_pose"
+                    object_name=config.task_config.pickup_obj_name, uuid="obj_start"
                 ),
                 GraspStateSensor(
                     object_name=config.task_config.pickup_obj_name,

molmo_spaces/tasks/pick_task.py

@@ -1,18 +1,18 @@
 import logging
 from typing import Any
 
+import gymnasium.spaces as gyms
 import numpy as np
 from scipy.spatial.transform import Rotation as R
-import gymnasium.spaces as gyms
 
 from molmo_spaces.configs.abstract_exp_config import MlSpacesExpConfig
 from molmo_spaces.configs.task_configs import PickTaskConfig
 from molmo_spaces.env.abstract_sensors import Sensor, SensorSuite
 from molmo_spaces.env.data_views import MlSpacesObject
 from molmo_spaces.env.sensors import (
-    get_core_sensors,
     GraspStateSensor,
     ObjectStartPoseSensor,
+    get_core_sensors,
 )
 from molmo_spaces.tasks.task import BaseMujocoTask
 from molmo_spaces.utils.mj_model_and_data_utils import descendant_geoms
@@ -60,14 +60,14 @@ def _create_sensor_suite_from_config(self, config: MlSpacesExpConfig) -> SensorS
         sensors.extend(
             [
                 ObjectStartPoseSensor(
-                    object_name=config.task_config.pickup_obj_name, uuid="obj_start_pose"
+                    object_name=config.task_config.pickup_obj_name, uuid="obj_start"
                 ),
                 GraspStateSensor(
                     object_name=config.task_config.pickup_obj_name,
                     uuid="grasp_state_pickup_obj",
                 ),
                 PickupObjGoalPoseSensor(
-                    uuid="obj_end_pose",
+                    uuid="obj_end",
                 ),
             ]
         )

molmo_spaces/utils/save_utils.py

@@ -732,39 +732,6 @@ def _save_extra_data_from_batched(obs_group, episode_data) -> None:
     """Save extra task data (pose sensors) from batched observations."""
     extra_group = obs_group.create_group("extra")
 
-    # TODO(max): why do we have this???
-    extra_sensor_mapping = {
-        # Standard object pose sensors
-        "obj_start_pose": "obj_start",
-        "obj_end_pose": "obj_end",
-        "grasp_state_pickup_obj": "grasp_state_pickup_obj",
-        "grasp_state_place_receptacle": "grasp_state_place_receptacle",
-        # Task info sensor
-        "task_info": "task_info",
-        # RBY1 door opening pose sensors
-        "door_start_pose": "obj_start",
-        "door_end_pose": "obj_end",
-        # RBY1 door state sensors
-        "door_state": "door_state",
-        "door_state_dict": "door_state_dict",
-        # Single arm TCP sensors
-        "tcp_pose": "tcp_pose",
-        "grasp_pose": "grasp_pose",
-        # RBY1 dual-arm TCP sensors
-        "left_tcp_pose": "left_tcp_pose",
-        "right_tcp_pose": "right_tcp_pose",
-        # RBY1 grasp state sensors
-        "rby1_left_grasp_state": "rby1_left_grasp_state",
-        "rby1_right_grasp_state": "rby1_right_grasp_state",
-        # Base pose sensor
-        "robot_base_pose": "robot_base_pose",
-        # Policy sensors
-        "policy_phase": "policy_phase",
-        "policy_num_retries": "policy_num_retries",
-        # Object tracking sensors
-        "object_image_points": "object_image_points",
-    }
-
     def _save_nested_data(data, group, name_prefix=""):
         """Recursively save nested dictionary data until hitting tensors."""
         if isinstance(data, dict):
@@ -790,11 +757,10 @@ def _save_nested_data(data, group, name_prefix=""):
                 except Exception as e:
                     log.warning(f"Could not save data for {name_prefix}: {type(data)}, error: {e}")
 
-    for sensor_name, target_name in extra_sensor_mapping.items():
-        if sensor_name in episode_data:
-            # Use recursive loop for all sensors - handles both simple tensors and nested dicts
-            sensor_data = episode_data[sensor_name]
-            _save_nested_data(sensor_data, extra_group, target_name)
+    for sensor_name in episode_data:
+        # Use recursive loop for all sensors - handles both simple tensors and nested dicts
+        sensor_data = episode_data[sensor_name]
+        _save_nested_data(sensor_data, extra_group, sensor_name)
 
 
 def _save_sensor_params_from_batched(obs_group, episode_data) -> None:

pyproject.toml

@@ -40,7 +40,7 @@ dependencies = [
     "msgpack-numpy",
     "mujoco-warp~=3.5.0",
     "mujoco-mjx~=3.5.0",
-    "warp-lang<1.15.0", # warp-lang had breaking changes in 1.15.0 that caused ik tests to fail
+    "warp-lang",
     "nbstripout>=0.7.0",
     "nltk~=3.9.2",
     "numpy>=1.19.0,<3",