diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5362b70b..d410efa5 100644 --- a/.pre-commit-config.yaml +++ b/.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"] diff --git a/docs/sensors.md b/docs/sensors.md index bb64f58a..907598da 100644 --- a/docs/sensors.md +++ b/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 | diff --git a/molmo_spaces/kinematics/parallel/warp_kinematics.py b/molmo_spaces/kinematics/parallel/warp_kinematics.py index 3b352482..42a55f39 100644 --- a/molmo_spaces/kinematics/parallel/warp_kinematics.py +++ b/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,7 +177,7 @@ 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] @@ -185,7 +185,7 @@ def cholesky_solve6(H: mat66f, b: vec6f) -> vec6f: # 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"): diff --git a/molmo_spaces/tasks/pick_and_place_task.py b/molmo_spaces/tasks/pick_and_place_task.py index 8cf9daf8..366de4d4 100644 --- a/molmo_spaces/tasks/pick_and_place_task.py +++ b/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, diff --git a/molmo_spaces/tasks/pick_task.py b/molmo_spaces/tasks/pick_task.py index ebbc59ec..2b422aae 100644 --- a/molmo_spaces/tasks/pick_task.py +++ b/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", ), ] ) diff --git a/molmo_spaces/utils/save_utils.py b/molmo_spaces/utils/save_utils.py index 1a46eabb..6f07d017 100644 --- a/molmo_spaces/utils/save_utils.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 68477b23..627ad1cb 100644 --- a/pyproject.toml +++ b/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",