diff --git a/README.md b/README.md index 4751229..3db8e49 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,27 @@ from orca_sim import OrcaHandCombinedExtended env = OrcaHandCombinedExtended(version="v1") # loads the v1 hand ``` -See our [`random_policy.py`](random_policy.py) example to see how to instantiate and interface an ORCA hand. +See our [`random_policy.py`](random_policy.py) example to see how to instantiate and interface an ORCA hand through the Gymnasium API. +If you want a lower-level demo built directly on the shared `orca_core` hand helpers, see [`hand_demo.py`](hand_demo.py). + +## Sample task: cube stacking + +Two cubes on a table, available with no robot, OrcaArm, or OrcaPanda: + +```python +from orca_sim import CubeStackingTabletop, OrcaArmCubeStacking, OrcaPandaCubeStacking + +env = OrcaArmCubeStacking() +obs, info = env.reset(seed=0) +obs, reward, terminated, truncated, info = env.step(env.action_space.sample()) +``` +Or use the OrcaPanda robot in the same way: + +```python +env = OrcaPandaCubeStacking() +obs, info = env.reset(seed=0) +# ... +``` ## Sample task: in-hand cube orientation @@ -112,4 +132,3 @@ The implementation is intentionally split so it doubles as a porting template: - The task scene lives in [`src/orca_sim/scenes/v2/scene_right_cube_orientation.xml`](src/orca_sim/scenes/v2/scene_right_cube_orientation.xml) and composes the existing hand MJCF with a single task cube. - The nominal palm-up hand pose and in-palm cube spawn are now authored into the task-specific scene/model files, so opening the XML directly in MuJoCo shows the intended setup. - The task logic lives in [`src/orca_sim/task_envs.py`](src/orca_sim/task_envs.py), including reset-time cube randomization and optional hand-pose overrides for custom MJCF layouts. - diff --git a/hand_demo.py b/hand_demo.py new file mode 100644 index 0000000..bef34e6 --- /dev/null +++ b/hand_demo.py @@ -0,0 +1,143 @@ +import argparse +import logging + +from orca_core.demo_presets import DEMO_POSE_FRACTIONS, DEMO_SEQUENCES + +from orca_sim import SimOrcaHand + + +SCENE_CHOICES = ( + "scene_left.xml", + "scene_right.xml", + "scene_combined.xml", + "scene_left_extended.xml", + "scene_right_extended.xml", + "scene_combined_extended.xml", + "scene_right_cube_orientation.xml", +) + + +def _expand_fractions_bimanual(fractions: dict[str, float]) -> dict[str, float]: + """Prefix bare canonical joint fractions with both ``left_`` and ``right_``.""" + return { + f"{side}_{joint}": value + for side in ("left", "right") + for joint, value in fractions.items() + } + + +def run_demo( + hand: SimOrcaHand, + demo_name: str = "main", + cycles: int = 1, + num_steps: int = 25, + step_size: float = 0.05, + return_to_neutral: bool = True, +) -> None: + """Run a named demo sequence on any scene, including combined (bimanual) ones. + + For single-hand scenes this delegates directly to ``hand.run_demo``. For + combined scenes the bare canonical joint names in the demo presets are + automatically expanded to ``left_*`` / ``right_*`` so both hands move in + sync. + """ + is_bimanual = hand.config.type is None # bimanual hand has no "left" or "right"type + if not is_bimanual: + return hand.run_demo( + demo_name=demo_name, + cycles=cycles, + num_steps=num_steps, + step_size=step_size, + return_to_neutral=return_to_neutral, + ) + + if demo_name not in DEMO_POSE_FRACTIONS: + available = ", ".join(sorted(DEMO_SEQUENCES)) + raise ValueError(f"Unknown demo '{demo_name}'. Available demos: {available}.") + + for name, fractions in DEMO_POSE_FRACTIONS[demo_name].items(): + hand.register_position( + name, + hand.pose_from_fractions(_expand_fractions_bimanual(fractions)), + ) + + hand.play_named_positions( + DEMO_SEQUENCES[demo_name], + cycles=cycles, + num_steps=num_steps, + step_size=step_size, + return_to_neutral=return_to_neutral, + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Play back orca_core demo presets on SimOrcaHand." + ) + parser.add_argument( + "--scene-file", + choices=SCENE_CHOICES, + default="scene_right.xml", + help="Scene XML to load.", + ) + parser.add_argument( + "--version", + default=None, + help="Embodiment version to load, e.g. 'v1' or 'v2'.", + ) + parser.add_argument( + "--render-mode", + choices=["human", "rgb_array"], + default="human", + help="Use 'human' for the MuJoCo viewer or 'rgb_array' for offscreen rendering.", + ) + parser.add_argument( + "--demo-name", + choices=list(DEMO_SEQUENCES), + default="main", + help="Which demo sequence to play.", + ) + parser.add_argument( + "--cycles", + type=int, + default=10, + help="Number of times to repeat the demo sequence.", + ) + parser.add_argument( + "--num-steps", + type=int, + default=25, + help="Interpolation steps between poses. More steps = smoother motion.", + ) + parser.add_argument( + "--step-size", + type=float, + default=0.05, + help="Sleep in seconds between interpolation steps. Increase to slow down playback.", + ) + args = parser.parse_args() + + hand = SimOrcaHand( + scene_file=args.scene_file, + version=args.version, + render_mode=args.render_mode, + ) + logging.info(f"scene={args.scene_file} version={hand.version} num_joints={len(hand.config.joint_ids)}") + + try: + hand.reset() + run_demo( + hand, + demo_name=args.demo_name, + cycles=args.cycles, + num_steps=args.num_steps, + step_size=args.step_size, + ) + except KeyboardInterrupt: + logging.info("Demo stopped by user.") + finally: + hand.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index a17641c..ebb6f40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,16 +15,35 @@ dependencies = [ "gymnasium>=0.29", "mujoco>=3.1", "numpy>=1.26", + "orca_core @ git+https://github.com/orcahand/orca_core.git@91ef976cdf8206d2b937f765b1d1bce091b9e2c5", ] [project.optional-dependencies] dev = [ "pytest>=8", + "pytest-mock>=3.14", +] +arm = [ + "mani_skill @ git+https://github.com/fracapuano/ManiSkill.git@3dc8c0c52c71c2f41932eb9f0371ca340cf1b038", + "torch", + "sapien", + "transforms3d", + "imageio", + "imageio-ffmpeg", + "orca_arm", + "pin" +] + +lerobot = [ + "lerobot" ] [tool.pytest.ini_options] testpaths = ["tests"] addopts = "-ra" +markers = [ + "slow: long-running test intended primarily for CI", +] [tool.setuptools] package-dir = {"" = "src"} @@ -36,6 +55,7 @@ where = ["src"] [tool.setuptools.package-data] orca_sim = [ "*.xml", + "scenes/*.xml", "scenes/*/*.xml", "assets/mjcf/*.xml", "assets/mjcf/*/*.xml", @@ -45,5 +65,9 @@ orca_sim = [ "assets/mjcf/right/collision/*.stl", "models/mjcf/*.mjcf", "models/*/*.mjcf", + "models/*/assets/*.xml", + "models/*/assets/*/*.stl", + "models/*/mjcf/*.xml", + "models/*/mjcf/*.mjcf", "models/*/assets/mjcf/*/*.stl", ] diff --git a/src/orca_sim/__init__.py b/src/orca_sim/__init__.py index 6c6e292..d739060 100644 --- a/src/orca_sim/__init__.py +++ b/src/orca_sim/__init__.py @@ -2,18 +2,44 @@ latest_version, list_versions, ) -from orca_sim.envs import ( - OrcaHandCombined, - OrcaHandCombinedExtended, - OrcaHandLeft, - OrcaHandLeftExtended, - OrcaHandRight, - OrcaHandRightExtended, -) from orca_sim.registry import register_envs -from orca_sim.task_envs import OrcaHandRightCubeOrientation + +try: + from orca_sim.envs import ( + OrcaHandCombined, + OrcaHandCombinedExtended, + OrcaHandLeft, + OrcaHandLeftExtended, + OrcaHandRight, + OrcaHandRightExtended, + ) + from orca_sim.hand import SimOrcaHand, SimOrcaHandConfig + from orca_sim.task_envs import ( + CubeStackingTabletop, + OrcaArmCubeStacking, + OrcaHandRightCubeOrientation, + OrcaPandaCubeStacking, + ) +except ModuleNotFoundError as exc: + if exc.name != "mujoco": + raise + + CubeStackingTabletop = None + OrcaArmCubeStacking = None + OrcaHandCombined = None + OrcaHandCombinedExtended = None + OrcaHandLeft = None + OrcaHandLeftExtended = None + OrcaHandRight = None + OrcaHandRightCubeOrientation = None + OrcaHandRightExtended = None + OrcaPandaCubeStacking = None + SimOrcaHand = None + SimOrcaHandConfig = None __all__ = [ + "CubeStackingTabletop", + "OrcaArmCubeStacking", "OrcaHandCombined", "OrcaHandCombinedExtended", "OrcaHandLeft", @@ -21,6 +47,9 @@ "OrcaHandRight", "OrcaHandRightCubeOrientation", "OrcaHandRightExtended", + "OrcaPandaCubeStacking", + "SimOrcaHand", + "SimOrcaHandConfig", "latest_version", "list_versions", "register_envs", diff --git a/src/orca_sim/builders/__init__.py b/src/orca_sim/builders/__init__.py new file mode 100644 index 0000000..6db0016 --- /dev/null +++ b/src/orca_sim/builders/__init__.py @@ -0,0 +1,11 @@ +from orca_sim.builders.orcaarm_camera_mjcf import ( + build_orcaarm_camera_mjcf, + camera_names, +) +from orca_sim.builders.orcapanda_mjcf import build_orcapanda_mjcf + +__all__ = [ + "build_orcaarm_camera_mjcf", + "build_orcapanda_mjcf", + "camera_names", +] diff --git a/src/orca_sim/builders/build_orcaarm_camera_mjcf.py b/src/orca_sim/builders/build_orcaarm_camera_mjcf.py new file mode 100644 index 0000000..66e1bed --- /dev/null +++ b/src/orca_sim/builders/build_orcaarm_camera_mjcf.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +from orca_sim.builders.orcaarm_camera_mjcf import build_orcaarm_camera_mjcf + + +def default_output_path() -> Path: + return ( + Path(__file__).resolve().parents[1] + / "scenes" + / "includes" + / "orcabot_with_cameras.xml" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--debug-sites", + action="store_true", + help="Add visible marker sites for camera mounts and aim targets.", + ) + parser.add_argument( + "--output", + type=Path, + default=default_output_path(), + help="Path to write the built MJCF include.", + ) + args = parser.parse_args() + + print(build_orcaarm_camera_mjcf(args.output, debug_sites=args.debug_sites)) + + +if __name__ == "__main__": + main() diff --git a/src/orca_sim/builders/build_orcapanda_mjcf.py b/src/orca_sim/builders/build_orcapanda_mjcf.py new file mode 100644 index 0000000..8c36911 --- /dev/null +++ b/src/orca_sim/builders/build_orcapanda_mjcf.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from pathlib import Path + +from orca_sim.builders.orcapanda_mjcf import build_orcapanda_mjcf + + +def default_output_path() -> Path: + return ( + Path(__file__).resolve().parents[1] + / "scenes" + / "includes" + / "orcapanda_namespaced.xml" + ) + + +def main() -> None: + print(build_orcapanda_mjcf(default_output_path())) + + +if __name__ == "__main__": + main() diff --git a/src/orca_sim/builders/orcaarm_camera_mjcf.py b/src/orca_sim/builders/orcaarm_camera_mjcf.py new file mode 100644 index 0000000..9f5b6a2 --- /dev/null +++ b/src/orca_sim/builders/orcaarm_camera_mjcf.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from pathlib import Path +import tempfile +import xml.etree.ElementTree as ET + + +CAMERA_SPECS = { + "openarm_body_link0": { + "name": "chest_table_camera", + "pos": "0.08 0 0.72", + "mode": "targetbody", + "target": "table", + "fovy": "75", + }, + "orcahand_left_ForeArmStructure-Model_e18f2368": { + "name": "left_wrist_camera", + "pos": "-0.01 0.085 -0.04", + "mode": "targetbody", + "target": "left_wrist_camera_target", + "fovy": "85", + }, + "orcahand_right_ForeArmStructure-Model_e18f2368": { + "name": "right_wrist_camera", + "pos": "-0.01 0.085 -0.04", + "mode": "targetbody", + "target": "right_wrist_camera_target", + "fovy": "85", + }, +} + +CAMERA_TARGET_SPECS = { + "orcahand_left_ForeArmStructure-Model_e18f2368": { + "name": "left_wrist_camera_target", + "pos": "-0.01 0.175 -0.075", + "site": { + "name": "left_wrist_camera_target_site", + "type": "sphere", + "pos": "0 0 0", + "size": "0.006", + "rgba": "1 0.85 0.05 1", + "group": "3", + }, + }, + "orcahand_right_ForeArmStructure-Model_e18f2368": { + "name": "right_wrist_camera_target", + "pos": "-0.01 0.175 -0.075", + "site": { + "name": "right_wrist_camera_target_site", + "type": "sphere", + "pos": "0 0 0", + "size": "0.006", + "rgba": "1 0.85 0.05 1", + "group": "3", + }, + }, +} + +CAMERA_SITE_SPECS = { + "chest_table_camera": { + "name": "chest_table_camera_site", + "type": "sphere", + "pos": CAMERA_SPECS["openarm_body_link0"]["pos"], + "size": "0.015", + "rgba": "0.1 1 0.1 1", + "group": "3", + }, + "left_wrist_camera": { + "name": "left_wrist_camera_site", + "type": "sphere", + "pos": CAMERA_SPECS["orcahand_left_ForeArmStructure-Model_e18f2368"]["pos"], + "size": "0.01", + "rgba": "1 0.1 0.1 1", + "group": "3", + }, + "right_wrist_camera": { + "name": "right_wrist_camera_site", + "type": "sphere", + "pos": CAMERA_SPECS["orcahand_right_ForeArmStructure-Model_e18f2368"]["pos"], + "size": "0.01", + "rgba": "0.1 0.3 1 1", + "group": "3", + }, +} + + +def camera_names() -> tuple[str, ...]: + return tuple(spec["name"] for spec in CAMERA_SPECS.values()) + + +def build_orcaarm_camera_mjcf( + output_path: str | Path | None = None, + *, + debug_sites: bool = False, +) -> Path: + import orca_arm + + source_path = Path(orca_arm.MJCF_PATH).resolve() + if output_path is None: + output_path = ( + Path(tempfile.gettempdir()) + / "orca_sim" + / f"{source_path.stem}_with_cameras.xml" + ) + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + + tree = ET.parse(source_path) + root = tree.getroot() + for mesh in root.findall(".//mesh"): + mesh_file = mesh.get("file") + if mesh_file is None: + continue + mesh_path = Path(mesh_file) + if not mesh_path.is_absolute(): + mesh.set("file", str((source_path.parent / mesh_path).resolve())) + + for body_name, camera_spec in CAMERA_SPECS.items(): + body = root.find(f".//body[@name='{body_name}']") + if body is None: + raise RuntimeError( + f"Unable to mount camera on {body_name!r}; body is missing in {source_path}." + ) + for existing in body.findall("camera"): + if existing.get("name") == camera_spec["name"]: + body.remove(existing) + site_spec = CAMERA_SITE_SPECS[camera_spec["name"]] + for existing in body.findall("site"): + if existing.get("name") == site_spec["name"]: + body.remove(existing) + target_spec = CAMERA_TARGET_SPECS.get(body_name) + if target_spec is not None: + for existing in body.findall("body"): + if existing.get("name") == target_spec["name"]: + body.remove(existing) + target_body = ET.SubElement( + body, + "body", + {"name": target_spec["name"], "pos": target_spec["pos"]}, + ) + if debug_sites: + ET.SubElement(target_body, "site", target_spec["site"]) + ET.SubElement(body, "camera", camera_spec) + if debug_sites: + ET.SubElement(body, "site", site_spec) + + tree.write(output_path, encoding="unicode") + return output_path diff --git a/src/orca_sim/builders/orcapanda_mjcf.py b/src/orca_sim/builders/orcapanda_mjcf.py new file mode 100644 index 0000000..9165e59 --- /dev/null +++ b/src/orca_sim/builders/orcapanda_mjcf.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +from pathlib import Path +import tempfile +import xml.etree.ElementTree as ET + + +MOUNT_POS = "-0.05 0 0.08" +ASSET_PREFIX = "orcapanda_" +WRIST_CAMERA_BODY = "orcahand_right_ForeArmStructure-Model_e18f2368" +WRIST_CAMERA_SPEC = { + "name": "orcapanda_wrist_camera", + "pos": "-0.01 0.085 -0.04", + "mode": "targetbody", + "target": "orcapanda_wrist_camera_target", + "fovy": "85", +} +WRIST_CAMERA_TARGET_SPEC = { + "name": "orcapanda_wrist_camera_target", + "pos": "-0.01 0.175 -0.075", +} + + +def build_orcapanda_mjcf(output_path: str | Path | None = None) -> Path: + source_path = _resolve_orcapanda_mjcf_path() + if output_path is None: + output_path = ( + Path(tempfile.gettempdir()) + / "orca_sim" + / f"{source_path.stem}_namespaced.xml" + ) + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + + tree = ET.parse(source_path) + root = tree.getroot() + _strip_global_elements(root) + _namespace_default_classes(root) + _namespace_assets(root, source_path) + _mount_base(root) + _add_wrist_camera(root) + + tree.write(output_path, encoding="unicode") + return output_path + + +def _resolve_orcapanda_mjcf_path() -> Path: + try: + import orca_arm + except ModuleNotFoundError: + package_path = None + else: + package_path = getattr(orca_arm, "ORCAPANDA_MJCF_PATH", None) + if package_path is not None: + source_path = Path(package_path).resolve() + if source_path.exists(): + return source_path + + repo_path = ( + Path(__file__).resolve().parents[3].parent + / "orca_arm" + / "orca_arm" + / "orcapanda.xml" + ) + if repo_path.exists(): + return repo_path + + raise RuntimeError( + "Could not find orcapanda.xml. Reinstall the local orca_arm package with " + "`python -m pip install -e ../orca_arm`, or update the fallback path in " + "this builder." + ) + + +def _strip_global_elements(root: ET.Element) -> None: + for tag in ("compiler", "option", "keyframe"): + for element in root.findall(tag): + root.remove(element) + + +def _namespace_default_classes(root: ET.Element) -> None: + class_names = { + element.get("class") + for element in root.findall(".//default") + if element.get("class") is not None + } + class_name_map = { + class_name: f"{ASSET_PREFIX}{class_name}" for class_name in class_names + } + for element in root.iter(): + class_name = element.get("class") + if class_name in class_name_map: + element.set("class", class_name_map[class_name]) + childclass_name = element.get("childclass") + if childclass_name in class_name_map: + element.set("childclass", class_name_map[childclass_name]) + + +def _namespace_assets(root: ET.Element, source_path: Path) -> None: + material_name_map: dict[str, str] = {} + for material in root.findall(".//material"): + name = material.get("name") + if name is None: + continue + namespaced_name = f"{ASSET_PREFIX}{name}" + material_name_map[name] = namespaced_name + material.set("name", namespaced_name) + + mesh_name_map: dict[str, str] = {} + for mesh in root.findall(".//mesh"): + mesh_file = mesh.get("file") + if mesh_file is None: + continue + + mesh_path = Path(mesh_file) + if not mesh_path.is_absolute(): + mesh.set("file", str((source_path.parent / mesh_path).resolve())) + if mesh.get("scale") is None: + mesh.set("scale", "1 1 1") + + name = mesh.get("name") + if name is None: + name = Path(mesh_file).stem + namespaced_name = f"{ASSET_PREFIX}{name}" + mesh_name_map[name] = namespaced_name + mesh.set("name", namespaced_name) + + for element in root.iter(): + material = element.get("material") + if material in material_name_map: + element.set("material", material_name_map[material]) + + mesh = element.get("mesh") + if mesh in mesh_name_map: + element.set("mesh", mesh_name_map[mesh]) + + +def _mount_base(root: ET.Element) -> None: + base = root.find(".//body[@name='panda_link0']") + if base is None: + raise RuntimeError("Unable to mount OrcaPanda; body 'panda_link0' is missing.") + base.set("pos", MOUNT_POS) + + +def _add_wrist_camera(root: ET.Element) -> None: + body = root.find(f".//body[@name='{WRIST_CAMERA_BODY}']") + if body is None: + raise RuntimeError( + f"Unable to mount OrcaPanda wrist camera; body {WRIST_CAMERA_BODY!r} is missing." + ) + + for existing in body.findall("camera"): + if existing.get("name") == WRIST_CAMERA_SPEC["name"]: + body.remove(existing) + for existing in body.findall("body"): + if existing.get("name") == WRIST_CAMERA_TARGET_SPEC["name"]: + body.remove(existing) + + ET.SubElement(body, "body", WRIST_CAMERA_TARGET_SPEC) + ET.SubElement(body, "camera", WRIST_CAMERA_SPEC) diff --git a/src/orca_sim/envs.py b/src/orca_sim/envs.py index b03bd00..7033750 100644 --- a/src/orca_sim/envs.py +++ b/src/orca_sim/envs.py @@ -1,18 +1,15 @@ -import sys from typing import Any import gymnasium as gym -import mujoco import numpy as np from gymnasium import spaces -from orca_sim.versions import ( - resolve_scene_path, -) +from orca_sim.hand import SimOrcaHand +RENDER_FPS = 30 class BaseOrcaHandEnv(gym.Env[np.ndarray, np.ndarray]): - metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30} + metadata = {"render_modes": ["human", "rgb_array"], "render_fps": RENDER_FPS} def __init__( self, @@ -22,24 +19,20 @@ def __init__( render_mode: str | None = None, ) -> None: super().__init__() - if render_mode not in {None, "human", "rgb_array"}: - raise ValueError(f"Unsupported render_mode: {render_mode}") - - self.scene_path = resolve_scene_path(scene_file, version=version) - self.version = self.scene_path.parent.name - self.frame_skip = frame_skip - self.render_mode = render_mode - - self.model = mujoco.MjModel.from_xml_path(str(self.scene_path)) - self.data = mujoco.MjData(self.model) - - self._default_camera = "closeup" - self._renderer: mujoco.Renderer | None = None - self._viewer: Any | None = None - - ctrl_range = self.model.actuator_ctrlrange.copy() - self.action_low = ctrl_range[:, 0].astype(np.float32) - self.action_high = ctrl_range[:, 1].astype(np.float32) + self.hand = SimOrcaHand( + scene_file=scene_file, + version=version, + frame_skip=frame_skip, + render_mode=render_mode, + ) + self.scene_path = self.hand.scene_path + self.version = self.hand.version + self.frame_skip = self.hand.frame_skip + self.render_mode = self.hand.render_mode + self.model = self.hand.model + self.data = self.hand.data + self.action_low = self.hand.action_low + self.action_high = self.hand.action_high self.action_space = spaces.Box( low=self.action_low, high=self.action_high, @@ -55,7 +48,7 @@ def __init__( ) def _get_obs(self) -> np.ndarray: - return np.concatenate([self.data.qpos.copy(), self.data.qvel.copy()]) + return self.hand.observe() def _get_reward(self) -> float: return 0.0 @@ -76,40 +69,13 @@ def reset( options: dict[str, Any] | None = None, ) -> tuple[np.ndarray, dict[str, Any]]: super().reset(seed=seed) - mujoco.mj_resetData(self.model, self.data) - - if options and "qpos" in options: - qpos = np.asarray(options["qpos"], dtype=np.float64) - if qpos.shape != self.data.qpos.shape: - raise ValueError( - f"Expected qpos shape {self.data.qpos.shape}, got {qpos.shape}" - ) - self.data.qpos[:] = qpos - - if options and "qvel" in options: - qvel = np.asarray(options["qvel"], dtype=np.float64) - if qvel.shape != self.data.qvel.shape: - raise ValueError( - f"Expected qvel shape {self.data.qvel.shape}, got {qvel.shape}" - ) - self.data.qvel[:] = qvel - - mujoco.mj_forward(self.model, self.data) - - if self.render_mode == "human": - self.render() + options = {} if options is None else dict(options) + self.hand.reset(qpos=options.get("qpos"), qvel=options.get("qvel")) return self._get_obs(), self._get_info() def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]: - action = np.asarray(action, dtype=np.float32) - if action.shape != self.action_space.shape: - raise ValueError( - f"Expected action shape {self.action_space.shape}, got {action.shape}" - ) - - self.data.ctrl[:] = np.clip(action, self.action_low, self.action_high) - mujoco.mj_step(self.model, self.data, nstep=self.frame_skip) + self.hand.step(action) obs = self._get_obs() reward = self._get_reward() @@ -117,49 +83,13 @@ def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, bool, dict[ truncated = self._get_truncated() info = self._get_info() - if self.render_mode == "human": - self.render() - return obs, reward, terminated, truncated, info def render(self) -> np.ndarray | None: - if self.render_mode == "rgb_array": - if self._renderer is None: - self._renderer = mujoco.Renderer(self.model) - self._renderer.update_scene(self.data) - return self._renderer.render() - - if self.render_mode == "human": - if self._viewer is None: - from mujoco import viewer - - try: - self._viewer = viewer.launch_passive(self.model, self.data) - except RuntimeError as exc: - if sys.platform == "darwin" and "mjpython" in str(exc): - raise RuntimeError( - "On macOS, MuJoCo human rendering must be launched with " - "`mjpython`, not plain `python3`. Run " - "`mjpython scripts/smoke_test_env.py --render-mode human` " - "for the interactive viewer, or use " - "`python3 scripts/smoke_test_env.py --render-mode rgb_array` " - "for an offscreen smoke test." - ) from exc - raise - # Force viewer free-camera to scene.xml defaults on startup. - mujoco.mjv_defaultFreeCamera(self.model, self._viewer.cam) - self._viewer.sync() - return None - - return None + return self.hand.render() def close(self) -> None: - if self._renderer is not None: - self._renderer.close() - self._renderer = None - if self._viewer is not None: - self._viewer.close() - self._viewer = None + self.hand.close() class OrcaHandLeft(BaseOrcaHandEnv): diff --git a/src/orca_sim/hand.py b/src/orca_sim/hand.py new file mode 100644 index 0000000..f1df6f9 --- /dev/null +++ b/src/orca_sim/hand.py @@ -0,0 +1,375 @@ +from __future__ import annotations + +import sys +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import mujoco +import numpy as np +from orca_core.base_hand import BaseHand +from orca_core.hand_config import BaseHandConfig +from orca_core.joint_position import OrcaJointPositions + +from orca_sim.joint_mapping import resolve_joint_mapping +from orca_sim.versions import resolve_scene_path + + +@dataclass(frozen=True, kw_only=True) +class SimOrcaHandConfig(BaseHandConfig): + scene_file: str + scene_path: str + version: str + scene_joint_names: tuple[str, ...] + actuator_ids: tuple[int, ...] + actuator_qpos_indices: tuple[int, ...] + actuator_qvel_indices: tuple[int, ...] + action_low: tuple[float, ...] + action_high: tuple[float, ...] + + @classmethod + def from_config_path( + cls, + config_path: str | None = None, + *, + scene_file: str = "scene_right.xml", + version: str | None = None, + joint_name_to_scene_joint_name: Mapping[str, str] | None = None, + hand_type: str | None = None, + model_version: str | None = None, + model_name: str | None = None, + ) -> "SimOrcaHandConfig": + del model_name + if version is None: + version = model_version + + if config_path is None: + scene_path = resolve_scene_path(scene_file, version=version) + else: + scene_path = Path(config_path).expanduser().resolve() + if not scene_path.exists(): + raise FileNotFoundError(f"Scene file not found: {scene_path}") + scene_file = scene_path.name + + model = mujoco.MjModel.from_xml_path(str(scene_path)) + + resolved_hand_type, resolved_mapping = resolve_joint_mapping( + scene_file=scene_file, + version=scene_path.parent.name, + joint_name_to_scene_joint_name=joint_name_to_scene_joint_name, + hand_type=hand_type, + ) + + actuator_metadata_by_scene_joint: dict[str, tuple[int, int, int, float, float]] = {} + for actuator_id in range(model.nu): + joint_id = int(model.actuator_trnid[actuator_id, 0]) + joint_name = model.joint(joint_id).name + if joint_name in actuator_metadata_by_scene_joint: + raise ValueError( + f"Scene joint {joint_name!r} is driven by multiple actuators, " + "which SimOrcaHandConfig does not currently support." + ) + + qpos_idx = int(model.jnt_qposadr[joint_id]) + qvel_idx = int(model.jnt_dofadr[joint_id]) + low, high = model.actuator_ctrlrange[actuator_id] + actuator_metadata_by_scene_joint[joint_name] = ( + actuator_id, + qpos_idx, + qvel_idx, + float(low), + float(high), + ) + + joint_ids: list[str] = [] + scene_joint_names: list[str] = [] + actuator_ids: list[int] = [] + joint_roms_dict: dict[str, list[float]] = {} + neutral_position: dict[str, float] = {} + actuator_qpos_indices: list[int] = [] + actuator_qvel_indices: list[int] = [] + action_low: list[float] = [] + action_high: list[float] = [] + + for joint_name, scene_joint_name in resolved_mapping.items(): + try: + actuator_id, qpos_idx, qvel_idx, low, high = actuator_metadata_by_scene_joint[ + scene_joint_name + ] + except KeyError as exc: + raise ValueError( + f"Scene joint {scene_joint_name!r} mapped from canonical joint " + f"{joint_name!r} is not actuator-controlled in {scene_path}." + ) from exc + + joint_ids.append(joint_name) + scene_joint_names.append(scene_joint_name) + actuator_ids.append(actuator_id) + joint_roms_dict[joint_name] = [low, high] + neutral_position[joint_name] = float(model.qpos0[qpos_idx]) + actuator_qpos_indices.append(qpos_idx) + actuator_qvel_indices.append(qvel_idx) + action_low.append(low) + action_high.append(high) + + return cls( + config_path=str(scene_path), + type=resolved_hand_type, + joint_ids=joint_ids, + joint_roms_dict=joint_roms_dict, + neutral_position=neutral_position, + scene_file=scene_file, + scene_path=str(scene_path), + version=scene_path.parent.name, + scene_joint_names=tuple(scene_joint_names), + actuator_ids=tuple(actuator_ids), + actuator_qpos_indices=tuple(actuator_qpos_indices), + actuator_qvel_indices=tuple(actuator_qvel_indices), + action_low=tuple(action_low), + action_high=tuple(action_high), + ) + + def validate(self) -> None: + super().validate() + + if not self.scene_file: + raise ValueError("scene_file must be provided for a simulated hand.") + if not self.scene_path: + raise ValueError("scene_path must be provided for a simulated hand.") + if not self.version: + raise ValueError("version must be provided for a simulated hand.") + + expected_len = len(self.joint_ids) + if len(self.scene_joint_names) != expected_len: + raise ValueError("Each canonical joint must map to exactly one scene joint.") + if len(self.actuator_ids) != expected_len: + raise ValueError("Each simulated joint must have an actuator id.") + if len(self.actuator_qpos_indices) != expected_len: + raise ValueError("Each simulated joint must have a qpos index.") + if len(self.actuator_qvel_indices) != expected_len: + raise ValueError("Each simulated joint must have a qvel index.") + if len(self.action_low) != expected_len: + raise ValueError("Each simulated joint must have a lower control bound.") + if len(self.action_high) != expected_len: + raise ValueError("Each simulated joint must have an upper control bound.") + + def __post_init__(self) -> None: + self.validate() + + @property + def joint_name_to_scene_joint_name(self) -> dict[str, str]: + return dict(zip(self.joint_ids, self.scene_joint_names, strict=True)) + + @property + def joint_name_to_actuator_id(self) -> dict[str, int]: + return dict(zip(self.joint_ids, self.actuator_ids, strict=True)) + + @property + def joint_name_to_qpos_idx(self) -> dict[str, int]: + return dict(zip(self.joint_ids, self.actuator_qpos_indices, strict=True)) + + @property + def joint_name_to_qvel_idx(self) -> dict[str, int]: + return dict(zip(self.joint_ids, self.actuator_qvel_indices, strict=True)) + + +class SimOrcaHand(BaseHand): + metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30} + config_cls = SimOrcaHandConfig + + def __init__( + self, + scene_file: str = "scene_right.xml", + version: str | None = None, + frame_skip: int = 5, + render_mode: str | None = None, + config_path: str | None = None, + config: SimOrcaHandConfig | None = None, + ) -> None: + if render_mode not in {None, "human", "rgb_array"}: + raise ValueError(f"Unsupported render_mode: {render_mode}") + + self.frame_skip = frame_skip + self.render_mode = render_mode + + super().__init__( + config_path=config_path, + config=config, + scene_file=scene_file, + version=version, + ) + + self.scene_path = Path(self.config.scene_path) + self.version = self.config.version + self.model = mujoco.MjModel.from_xml_path(str(self.scene_path)) + self.data = mujoco.MjData(self.model) + + self.action_low = np.asarray(self.config.action_low, dtype=np.float32) + self.action_high = np.asarray(self.config.action_high, dtype=np.float32) + + self._renderer: mujoco.Renderer | None = None + self._viewer: Any | None = None + + def _get_joint_positions(self) -> OrcaJointPositions: + return OrcaJointPositions.from_dict( + { + joint_name: float(self.data.qpos[self.config.joint_name_to_qpos_idx[joint_name]]) + for joint_name in self.config.joint_ids + } + ) + + def _set_joint_positions(self, joint_pos: OrcaJointPositions) -> bool: + current_joint_positions = self._get_joint_positions().as_dict() + for joint_name, value in joint_pos: + if joint_name not in current_joint_positions: + raise ValueError(f"Unknown canonical joint name: {joint_name}") + current_joint_positions[joint_name] = float(value) + + for joint_name, value in current_joint_positions.items(): + self.data.qpos[self.config.joint_name_to_qpos_idx[joint_name]] = value + # Setting a target joint configuration is a teleport-like state edit, + # so we explicitly zero the corresponding joint velocity. + self.data.qvel[self.config.joint_name_to_qvel_idx[joint_name]] = 0.0 + + ctrl = np.array( + [current_joint_positions[joint_name] for joint_name in self.config.joint_ids], + dtype=np.float32, + ) + self.set_control(ctrl) + + # step simulation forward based on current simulator's state + mujoco.mj_forward(self.model, self.data) + + if self.render_mode == "human": + self.render() + + return True + + def observe(self) -> np.ndarray: + return np.concatenate([self.data.qpos.copy(), self.data.qvel.copy()]) + + def _validate_qpos(self, qpos: np.ndarray) -> np.ndarray: + resolved_qpos = np.asarray(qpos, dtype=np.float64) + if resolved_qpos.shape != self.data.qpos.shape: + raise ValueError( + f"Expected qpos shape {self.data.qpos.shape}, got {resolved_qpos.shape}" + ) + return resolved_qpos + + def _validate_qvel(self, qvel: np.ndarray) -> np.ndarray: + resolved_qvel = np.asarray(qvel, dtype=np.float64) + if resolved_qvel.shape != self.data.qvel.shape: + raise ValueError( + f"Expected qvel shape {self.data.qvel.shape}, got {resolved_qvel.shape}" + ) + return resolved_qvel + + def _validate_ctrl(self, ctrl: np.ndarray) -> np.ndarray: + resolved_ctrl = np.asarray(ctrl, dtype=np.float32) + expected_shape = (len(self.config.joint_ids),) + if resolved_ctrl.shape != expected_shape: + raise ValueError(f"Expected action shape {expected_shape}, got {resolved_ctrl.shape}") + return np.clip(resolved_ctrl, self.action_low, self.action_high) + + def set_state( + self, + *, + qpos: np.ndarray | None = None, + qvel: np.ndarray | None = None, + ctrl: np.ndarray | None = None, + forward: bool = True, + ) -> np.ndarray: + if qpos is not None: + self.data.qpos[:] = self._validate_qpos(qpos) + + if qvel is not None: + self.data.qvel[:] = self._validate_qvel(qvel) + + if ctrl is not None: + self.set_control(ctrl) + + if forward: + self.forward() + + return self.observe() + + def set_control(self, ctrl: np.ndarray) -> np.ndarray: + resolved_ctrl = self._validate_ctrl(ctrl) + for actuator_id, value in zip( + self.config.actuator_ids, + resolved_ctrl, + strict=True, + ): + self.data.ctrl[actuator_id] = float(value) + return self.data.ctrl[list(self.config.actuator_ids)].copy() + + def forward(self) -> np.ndarray: + mujoco.mj_forward(self.model, self.data) + if self.render_mode == "human": + self.render() + + return self.observe() + + def reset( + self, + *, + qpos: np.ndarray | None = None, + qvel: np.ndarray | None = None, + ctrl: np.ndarray | None = None, + ) -> np.ndarray: + mujoco.mj_resetData(self.model, self.data) + default_ctrl = self._get_joint_positions().as_array(self.config.joint_ids).astype( + np.float32 + ) + return self.set_state( + qpos=qpos, + qvel=qvel, + ctrl=default_ctrl if ctrl is None else ctrl, + forward=True, + ) + + def step(self, action: np.ndarray | None = None, *, nstep: int | None = None) -> np.ndarray: + if action is not None: + self.set_control(action) # mujoco simulation is stateful - this updates the state for stepping + + mujoco.mj_step(self.model, self.data, nstep=self.frame_skip if nstep is None else nstep) + + if self.render_mode == "human": + self.render() + + return self.observe() + + def render(self) -> np.ndarray | None: + if self.render_mode == "rgb_array": + if self._renderer is None: + self._renderer = mujoco.Renderer(self.model) + self._renderer.update_scene(self.data) + return self._renderer.render() + + if self.render_mode == "human": + if self._viewer is None: + from mujoco import viewer + + try: + self._viewer = viewer.launch_passive(self.model, self.data) + except RuntimeError as exc: + if sys.platform == "darwin" and "mjpython" in str(exc): + raise RuntimeError( + "On macOS, MuJoCo human rendering must be launched with " + "`mjpython`, not plain `python3`." + ) from exc + raise + mujoco.mjv_defaultFreeCamera(self.model, self._viewer.cam) + self._viewer.sync() + return None + + return None + + def close(self) -> None: + if self._renderer is not None: + self._renderer.close() + self._renderer = None + if self._viewer is not None: + self._viewer.close() + self._viewer = None diff --git a/src/orca_sim/joint_mapping.py b/src/orca_sim/joint_mapping.py new file mode 100644 index 0000000..5bf3488 --- /dev/null +++ b/src/orca_sim/joint_mapping.py @@ -0,0 +1,111 @@ +from collections.abc import Mapping + +from orca_core import canonical_joint_ids + + +def canonical_single_hand_joint_ids( + *, + version: str | None = None, + hand_type: str | None = None, +) -> tuple[str, ...]: + try: + return canonical_joint_ids(version=version, type=hand_type) + except FileNotFoundError: + if hand_type == "left": + return canonical_joint_ids(version=version, type="right") + raise + + +V2_LOCAL_SCENE_JOINT_BY_CANONICAL: dict[str, str] = { + # NOTE: v2 uses a different thumb naming scheme in the MJCF + "wrist": "wrist", + "thumb_cmc": "t-cmc", + "thumb_abd": "t-abd", + "thumb_mcp": "t-mcp", + "thumb_dip": "t-pip", + "index_abd": "i-abd", + "index_mcp": "i-mcp", + "index_pip": "i-pip", + "middle_abd": "m-abd", + "middle_mcp": "m-mcp", + "middle_pip": "m-pip", + "ring_abd": "r-abd", + "ring_mcp": "r-mcp", + "ring_pip": "r-pip", + "pinky_abd": "p-abd", + "pinky_mcp": "p-mcp", + "pinky_pip": "p-pip", +} + + +def default_joint_name_to_scene_joint_name( + *, + scene_file: str, + version: str, +) -> tuple[str | None, dict[str, str]]: + if version == "v1": + canonical = canonical_single_hand_joint_ids(version=version, hand_type="right") + local_mapping = {joint: joint for joint in canonical} + elif version == "v2": + canonical = tuple(V2_LOCAL_SCENE_JOINT_BY_CANONICAL) + local_mapping = V2_LOCAL_SCENE_JOINT_BY_CANONICAL + else: + raise FileNotFoundError(f"Unsupported embodiment version for joint mapping: {version}") + + if "combined" in scene_file: + mapping: dict[str, str] = {} + for side in ("left", "right"): + for joint in canonical: + mapping[f"{side}_{joint}"] = f"{side}_{local_mapping[joint]}" + return None, mapping + + if "left" in scene_file: + if version == "v1": + canonical = canonical_single_hand_joint_ids(version=version, hand_type="left") + return ( + "left", + {joint: f"left_{local_mapping[joint]}" for joint in canonical}, + ) + + if "right" in scene_file: + canonical = canonical_single_hand_joint_ids(version=version, hand_type="right") + return ( + "right", + {joint: f"right_{local_mapping[joint]}" for joint in canonical}, + ) + + raise ValueError( + "Unable to infer a default hand-joint mapping from scene_file. " + "Provide joint_name_to_scene_joint_name explicitly." + ) + + +def infer_hand_type_from_joint_names(joint_names: list[str]) -> str | None: + prefixes = { + joint_name.split("_", maxsplit=1)[0] + for joint_name in joint_names + if "_" in joint_name + } + if prefixes == {"left"}: + return "left" + if prefixes == {"right"}: + return "right" + return None + + +def resolve_joint_mapping( + *, + scene_file: str, + version: str, + joint_name_to_scene_joint_name: Mapping[str, str] | None, + hand_type: str | None, +) -> tuple[str | None, dict[str, str]]: + if joint_name_to_scene_joint_name is None: + return default_joint_name_to_scene_joint_name( + scene_file=scene_file, + version=version, + ) + + resolved_mapping = dict(joint_name_to_scene_joint_name) + resolved_type = hand_type or infer_hand_type_from_joint_names(list(resolved_mapping)) + return resolved_type, resolved_mapping diff --git a/src/orca_sim/maniskill/__init__.py b/src/orca_sim/maniskill/__init__.py new file mode 100644 index 0000000..17eb233 --- /dev/null +++ b/src/orca_sim/maniskill/__init__.py @@ -0,0 +1,10 @@ +"""ManiSkill backend for orca_sim, used to create more advanced simulation environments +leveraging the Orca hand. + +Requires the ``maniskill`` optional dependency group: + pip install 'orca_sim[maniskill]' +""" +from .agent import OrcaArm # noqa: F401 + +# Example environment that uses the Orca hand. +from .push_cube_env import PushCubeOrcaArmEnv # noqa: F401 diff --git a/src/orca_sim/maniskill/agent.py b/src/orca_sim/maniskill/agent.py new file mode 100644 index 0000000..1d4b4aa --- /dev/null +++ b/src/orca_sim/maniskill/agent.py @@ -0,0 +1,190 @@ +"""ManiSkill agent definition for the orca_arm bimanual robot. + +Registers orca_arm under uid="orca_arm" with a single PD-joint-position +controller covering every movable joint in orca_arm's URDF. Picks the +right hand's wrist tower as the TCP for downstream task code. + +Joint groups are derived at import time from ``orca_arm.URDF_PATH`` by +splitting movable joint names by prefix, so adding/removing joints +upstream is reflected here automatically. + +Controller gains and force limits are read from ``orca_arm.MJCF_PATH`` so +the ManiSkill controller mirrors the MuJoCo position actuator defaults. +""" +import xml.etree.ElementTree as ET +from copy import deepcopy + +import numpy as np +import orca_arm +import sapien + +from mani_skill.agents.base_agent import BaseAgent, Keyframe +from mani_skill.agents.controllers import PDJointPosControllerConfig +from mani_skill.agents.registration import register_agent +from mani_skill.utils import sapien_utils + + +_JOINT_PREFIXES = ( + "openarm_left_", + "openarm_right_", + "orcahand_left_", + "orcahand_right_", +) + + +def _float_attr(element: ET.Element, attr: str, context: str) -> float: + value = element.get(attr) + if value is None: + raise RuntimeError(f"Expected {context} to define {attr!r}.") + return float(value) + + +def _symmetric_force_limit(position: ET.Element, context: str) -> float: + value = position.get("forcerange") + if value is None: + raise RuntimeError(f"Expected {context} to define 'forcerange'.") + + low, high = (float(v) for v in value.split()) + if not np.isclose(abs(low), abs(high)): + raise RuntimeError( + f"Expected {context} forcerange to be symmetric, got {value!r}." + ) + return abs(high) + + +def _pd_position_defaults(mjcf_path: str, default_class: str) -> dict[str, float]: + """Read ManiSkill PD position controller values from MJCF defaults.""" + context = f"orca_arm MJCF default class {default_class!r}" + default = ET.parse(mjcf_path).getroot().find( + f".//default[@class='{default_class}']" + ) + if default is None: + raise RuntimeError(f"Expected {context} to exist.") + + position = default.find("position") + if position is None: + raise RuntimeError(f"Expected {context} to define a position actuator.") + + return { + "stiffness": _float_attr(position, "kp", context), + "damping": _float_attr(position, "kv", context), + "force_limit": _symmetric_force_limit(position, context), + } + + +def _group_movable_joints(urdf_path: str) -> dict[str, list[str]]: + """Group movable joints in the URDF by name prefix, preserving URDF order. + + Order within each group = order of appearance in the URDF, which sets the + layout of the controller's action vector. Fails loudly if a movable joint + doesn't match any known prefix, so upstream additions can't silently drop + out of the action space. + """ + groups: dict[str, list[str]] = {p: [] for p in _JOINT_PREFIXES} + for joint in ET.parse(urdf_path).getroot().findall("joint"): + if joint.get("type") == "fixed": + continue + name = joint.get("name") + for prefix in _JOINT_PREFIXES: + if name.startswith(prefix): + groups[prefix].append(name) + break + else: + raise RuntimeError( + f"orca_arm URDF joint {name!r} matches no known prefix " + f"{_JOINT_PREFIXES}; update orca_sim agent groupings." + ) + return groups + + +_GROUPS = _group_movable_joints(orca_arm.URDF_PATH) +LEFT_ARM_JOINTS = _GROUPS["openarm_left_"] +RIGHT_ARM_JOINTS = _GROUPS["openarm_right_"] +LEFT_HAND_JOINTS = _GROUPS["orcahand_left_"] +RIGHT_HAND_JOINTS = _GROUPS["orcahand_right_"] + +ALL_ARM_JOINTS = LEFT_ARM_JOINTS + RIGHT_ARM_JOINTS +ALL_HAND_JOINTS = LEFT_HAND_JOINTS + RIGHT_HAND_JOINTS +ALL_JOINTS = ALL_ARM_JOINTS + ALL_HAND_JOINTS + +ARM_CONTROLLER_DEFAULTS = _pd_position_defaults(orca_arm.MJCF_PATH, "arm_joint") +HAND_CONTROLLER_DEFAULTS = _pd_position_defaults(orca_arm.MJCF_PATH, "hand_joint") + + +def _carpals_link(urdf_path: str, prefix: str) -> str: + """The unique ``Carpals`` link for a given hand — the wrist/palm body.""" + matches = [ + link.get("name") + for link in ET.parse(urdf_path).getroot().findall("link") + if link.get("name", "").startswith(prefix) and "Carpals" in link.get("name", "") + ] + if len(matches) != 1: + raise RuntimeError( + f"Expected exactly one 'Carpals' link with prefix {prefix!r}, found {matches}." + ) + return matches[0] + + +# Wrist link of the right hand — used as the TCP for tasks that need a +# single end-effector reference point. +RIGHT_TCP_LINK = _carpals_link(orca_arm.URDF_PATH, "orcahand_right_") +LEFT_TCP_LINK = _carpals_link(orca_arm.URDF_PATH, "orcahand_left_") + +@register_agent() +class OrcaArm(BaseAgent): + uid = "orca_arm" + urdf_path = orca_arm.URDF_PATH + + # NOTE: need an update to the OrcaArm with updated STLs + # to remove this disabling of self-collisions (see v0.0.0-release + # notes of OrcaArm) + disable_self_collisions = True + + arm_stiffness = ARM_CONTROLLER_DEFAULTS["stiffness"] + arm_damping = ARM_CONTROLLER_DEFAULTS["damping"] + arm_force_limit = ARM_CONTROLLER_DEFAULTS["force_limit"] + + hand_stiffness = HAND_CONTROLLER_DEFAULTS["stiffness"] + hand_damping = HAND_CONTROLLER_DEFAULTS["damping"] + hand_force_limit = HAND_CONTROLLER_DEFAULTS["force_limit"] + + keyframes = dict( + rest=Keyframe( + qpos=np.zeros(len(ALL_JOINTS)), + pose=sapien.Pose(), + ) + ) + + @property + def _controller_configs(self): + arm = PDJointPosControllerConfig( + ALL_ARM_JOINTS, + lower=None, + upper=None, + stiffness=self.arm_stiffness, + damping=self.arm_damping, + force_limit=self.arm_force_limit, + normalize_action=False, + ) + hand = PDJointPosControllerConfig( + ALL_HAND_JOINTS, + lower=None, + upper=None, + stiffness=self.hand_stiffness, + damping=self.hand_damping, + force_limit=self.hand_force_limit, + normalize_action=False, + ) + configs = dict( + pd_joint_pos=dict(arm=arm, hand=hand), + ) + return deepcopy(configs) + + def _after_init(self): + self.tcp = sapien_utils.get_obj_by_name( + self.robot.get_links(), RIGHT_TCP_LINK + ) + + @property + def tcp_pose(self): + return self.tcp.pose diff --git a/src/orca_sim/maniskill/push_cube_env.py b/src/orca_sim/maniskill/push_cube_env.py new file mode 100644 index 0000000..338e96c --- /dev/null +++ b/src/orca_sim/maniskill/push_cube_env.py @@ -0,0 +1,163 @@ +"""ManiSkill PushCube task variant using the orca_arm bimanual robot. + +Table + cube + goal target; same physics as ManiSkill's stock PushCube-v1, +but with orca_arm loaded in place of the panda. +""" +from typing import Any + +import numpy as np +import sapien +import torch +from transforms3d.euler import euler2quat + +from mani_skill.envs.sapien_env import BaseEnv +from mani_skill.sensors.camera import CameraConfig +from mani_skill.utils import sapien_utils +from mani_skill.utils.building import actors +from mani_skill.utils.registration import register_env +from mani_skill.utils.scene_builder.table import TableSceneBuilder +from mani_skill.utils.structs import Pose +from mani_skill.utils.structs.types import GPUMemoryConfig, SimConfig + +from orca_sim.maniskill.agent import OrcaArm + + +@register_env("PushCubeOrcaArm-v1", max_episode_steps=200) +class PushCubeOrcaArmEnv(BaseEnv): + """Push a cube into a goal region using the orca_arm bimanual robot.""" + + SUPPORTED_ROBOTS = ["orca_arm"] + agent: OrcaArm + + goal_radius = 0.1 + cube_half_size = 0.025 + + def __init__(self, *args, robot_uids="orca_arm", **kwargs): + super().__init__(*args, robot_uids=robot_uids, **kwargs) + + @property + def _default_sim_config(self): + return SimConfig( + gpu_memory_config=GPUMemoryConfig( + found_lost_pairs_capacity=2**25, + max_rigid_patch_count=2**18, + ) + ) + + @property + def _default_sensor_configs(self): + pose = sapien_utils.look_at(eye=[0.6, 0.0, 0.8], target=[0.0, 0.0, 0.2]) + return [ + CameraConfig( + "base_camera", + pose=pose, + width=128, + height=128, + fov=np.pi / 2, + near=0.01, + far=100, + ) + ] + + @property + def _default_human_render_camera_configs(self): + pose = sapien_utils.look_at([1.2, 1.0, 1.0], [0.0, 0.0, 0.4]) + return CameraConfig( + "render_camera", + pose=pose, + width=512, + height=512, + fov=1.0, + near=0.01, + far=100, + ) + + def _load_agent(self, options: dict): + # The orca_arm "world" link sits at z=0 in URDF coordinates and the + # base column extends upward, so we place it just behind the table + # edge so the hands hover over the table surface. + super()._load_agent(options, sapien.Pose(p=[-0.7, 0, 0])) + + def _load_scene(self, options: dict): + self.table_scene = TableSceneBuilder(env=self) + self.table_scene.build() + + self.obj = actors.build_cube( + self.scene, + half_size=self.cube_half_size, + color=np.array([12, 42, 160, 255]) / 255, + name="cube", + body_type="dynamic", + initial_pose=sapien.Pose(p=[0, 0, self.cube_half_size]), + ) + + self.goal_region = actors.build_red_white_target( + self.scene, + radius=self.goal_radius, + thickness=1e-5, + name="goal_region", + add_collision=False, + body_type="kinematic", + initial_pose=sapien.Pose(p=[0, 0, 1e-3]), + ) + + def _initialize_episode(self, env_idx: torch.Tensor, options: dict): + with torch.device(self.device): + b = len(env_idx) + self.table_scene.initialize(env_idx) + + xyz = torch.zeros((b, 3)) + xyz[..., :2] = torch.rand((b, 2)) * 0.2 - 0.1 + xyz[..., 2] = self.cube_half_size + self.obj.set_pose(Pose.create_from_pq(p=xyz, q=[1, 0, 0, 0])) + + target_xyz = xyz + torch.tensor([0.1 + self.goal_radius, 0, 0]) + target_xyz[..., 2] = 1e-3 + self.goal_region.set_pose( + Pose.create_from_pq( + p=target_xyz, + q=euler2quat(0, np.pi / 2, 0), + ) + ) + + def evaluate(self): + is_obj_placed = ( + torch.linalg.norm( + self.obj.pose.p[..., :2] - self.goal_region.pose.p[..., :2], axis=1 + ) + < self.goal_radius + ) & (self.obj.pose.p[..., 2] < self.cube_half_size + 5e-3) + return {"success": is_obj_placed} + + def _get_obs_extra(self, info: dict): + obs = dict(tcp_pose=self.agent.tcp_pose.raw_pose) + if self.obs_mode_struct.use_state: + obs.update( + goal_pos=self.goal_region.pose.p, + obj_pose=self.obj.pose.raw_pose, + ) + return obs + + def compute_dense_reward(self, obs: Any, action, info: dict): + # Distance from right-hand TCP to a push pose just behind the cube. + push_pose = Pose.create_from_pq( + p=self.obj.pose.p + + torch.tensor( + [-self.cube_half_size - 0.01, 0, 0], device=self.device + ) + ) + tcp_to_push = push_pose.p - self.agent.tcp_pose.p + reach_reward = 1 - torch.tanh(5 * torch.linalg.norm(tcp_to_push, axis=1)) + reward = reach_reward + + obj_to_goal = torch.linalg.norm( + self.obj.pose.p[..., :2] - self.goal_region.pose.p[..., :2], axis=1 + ) + place_reward = 1 - torch.tanh(5 * obj_to_goal) + reward += place_reward + + reward[info["success"]] = 4 + return reward + + def compute_normalized_dense_reward(self, obs: Any, action, info: dict): + return self.compute_dense_reward(obs, action, info) / 4.0 diff --git a/src/orca_sim/maniskill/run.py b/src/orca_sim/maniskill/run.py new file mode 100644 index 0000000..3c47f65 --- /dev/null +++ b/src/orca_sim/maniskill/run.py @@ -0,0 +1,86 @@ +"""Run the PushCubeOrcaArm-v1 environment with random actions. + +Usage: + python -m orca_sim.maniskill.run # interactive viewer + python -m orca_sim.maniskill.run --no-render # headless smoke test + python -m orca_sim.maniskill.run --record out.mp4 # offscreen video + +Requires the ``maniskill`` optional dependency group: + pip install 'orca_sim[maniskill]' +""" +import argparse + +import gymnasium as gym +import numpy as np + +# Importing the subpackage registers the agent and env with ManiSkill. +import orca_sim.maniskill # noqa: F401 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--no-render", action="store_true") + parser.add_argument( + "--record", + type=str, + default=None, + help="Path to save an mp4 of the rollout (rgb_array mode).", + ) + parser.add_argument("--steps", type=int, default=3000) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + if args.record or args.no_render: + render_mode = "rgb_array" + else: + render_mode = "human" + + env = gym.make( + "PushCubeOrcaArm-v1", + obs_mode="state", + control_mode="pd_joint_pos", + render_mode=render_mode, + ) + obs, info = env.reset(seed=args.seed) + print(f"action space: {env.action_space}") + print(f"obs keys: {list(obs.keys()) if isinstance(obs, dict) else type(obs)}") + + rng = np.random.default_rng(args.seed) + action_low = env.action_space.low + action_high = env.action_space.high + + frames = [] + for step in range(args.steps): + # Small jitter around the rest pose so the robot doesn't fly apart. + action = rng.uniform(action_low, action_high) * 0.05 + obs, reward, terminated, truncated, info = env.step(action) + if args.record: + frame = env.render() + if hasattr(frame, "cpu"): + frame = frame.cpu().numpy() + if frame.ndim == 4: + frame = frame[0] + frames.append(frame) + elif not args.no_render: + env.render() + if step % 30 == 0: + print( + f"step {step:4d} reward={float(reward):+.3f} " + f"success={bool(info.get('success', False))}" + ) + if terminated or truncated: + obs, info = env.reset() + + env.close() + + if args.record and frames: + import imageio.v2 as imageio + + imageio.mimsave(args.record, frames, fps=30) + print(f"saved {len(frames)} frames to {args.record}") + + print("done.") + + +if __name__ == "__main__": + main() diff --git a/src/orca_sim/registry.py b/src/orca_sim/registry.py index 8831415..60959e1 100644 --- a/src/orca_sim/registry.py +++ b/src/orca_sim/registry.py @@ -5,7 +5,7 @@ def register_envs() -> None: registry = gym.registry - specs = { + versioned_specs = { "OrcaHandLeft": ("orca_sim.envs:OrcaHandLeft", "scene_left.xml"), "OrcaHandLeftExtended": ( "orca_sim.envs:OrcaHandLeftExtended", @@ -26,8 +26,22 @@ def register_envs() -> None: "scene_right_cube_orientation.xml", ), } + unversioned_specs = { + "CubeStackingTabletop": ( + "orca_sim.task_envs:CubeStackingTabletop", + "cube_stacking.xml", + ), + "OrcaArmCubeStacking": ( + "orca_sim.task_envs:OrcaArmCubeStacking", + "orcaarm_cube_stacking.xml", + ), + "OrcaPandaCubeStacking": ( + "orca_sim.task_envs:OrcaPandaCubeStacking", + "orcapanda_cube_stacking.xml", + ), + } for version in list_versions(): - for env_name, (entry_point, scene_file) in specs.items(): + for env_name, (entry_point, scene_file) in versioned_specs.items(): try: resolve_scene_path(scene_file, version=version) except FileNotFoundError: @@ -40,3 +54,12 @@ def register_envs() -> None: entry_point=entry_point, kwargs={"version": version}, ) + + for env_id, (entry_point, scene_file) in unversioned_specs.items(): + try: + resolve_scene_path(scene_file) + except FileNotFoundError: + continue + + if env_id not in registry: + gym.register(id=env_id, entry_point=entry_point) diff --git a/src/orca_sim/scenes/cube_stacking.xml b/src/orca_sim/scenes/cube_stacking.xml new file mode 100644 index 0000000..1c280a3 --- /dev/null +++ b/src/orca_sim/scenes/cube_stacking.xml @@ -0,0 +1,57 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/orca_sim/scenes/includes/orcabot_with_cameras.xml b/src/orca_sim/scenes/includes/orcabot_with_cameras.xml new file mode 100644 index 0000000..c4fa29d --- /dev/null +++ b/src/orca_sim/scenes/includes/orcabot_with_cameras.xml @@ -0,0 +1,418 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/src/orca_sim/scenes/includes/orcapanda_namespaced.xml b/src/orca_sim/scenes/includes/orcapanda_namespaced.xml new file mode 100644 index 0000000..104f5ff --- /dev/null +++ b/src/orca_sim/scenes/includes/orcapanda_namespaced.xml @@ -0,0 +1,344 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/orca_sim/scenes/orcaarm_cube_stacking.xml b/src/orca_sim/scenes/orcaarm_cube_stacking.xml new file mode 100644 index 0000000..d2927f8 --- /dev/null +++ b/src/orca_sim/scenes/orcaarm_cube_stacking.xml @@ -0,0 +1,29 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/orca_sim/scenes/orcaarm_cube_stacking_cameras.xml b/src/orca_sim/scenes/orcaarm_cube_stacking_cameras.xml new file mode 100644 index 0000000..2b51818 --- /dev/null +++ b/src/orca_sim/scenes/orcaarm_cube_stacking_cameras.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/orca_sim/scenes/orcapanda_cube_stacking.xml b/src/orca_sim/scenes/orcapanda_cube_stacking.xml new file mode 100644 index 0000000..356e2a2 --- /dev/null +++ b/src/orca_sim/scenes/orcapanda_cube_stacking.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/orca_sim/task_envs.py b/src/orca_sim/task_envs.py index b8664a5..dc5fe69 100644 --- a/src/orca_sim/task_envs.py +++ b/src/orca_sim/task_envs.py @@ -1,6 +1,7 @@ from __future__ import annotations -from collections.abc import Mapping +import sys +from collections.abc import Mapping, Sequence from typing import Any import gymnasium as gym @@ -9,6 +10,688 @@ from gymnasium import spaces from orca_sim.envs import BaseOrcaHandEnv +from orca_sim.builders.orcaarm_camera_mjcf import ( + camera_names as default_orcaarm_camera_names, +) +from orca_sim.versions import SCENES_ROOT, resolve_scene_path + + +class CubeStackingTabletop(gym.Env[np.ndarray, np.ndarray]): + """Robot-free MuJoCo tabletop scene with two randomly placed stacking cubes.""" + + metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30} + CUBE_NAMES = ("red_cube", "blue_cube") + CUBE_JOINT_NAMES = ("red_cube_freejoint", "blue_cube_freejoint") + TARGET_BODY_NAME = "stack_target" + DEFAULT_WORKSPACE_BOUNDS = np.array( + [[0.38, 0.58], [-0.22, 0.22]], + dtype=np.float64, + ) + + def __init__( + self, + render_mode: str | None = None, + version: str | None = None, + *, + scene_file: str = "cube_stacking.xml", + frame_skip: int = 5, + workspace_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] = DEFAULT_WORKSPACE_BOUNDS, + min_cube_spacing: float = 0.09, + target_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] | None = None, + min_target_cube_spacing: float = 0.09, + randomize_yaw: bool = True, + render_camera: str = "topdown", + top_cube_name: str = "red_cube", + base_cube_name: str = "blue_cube", + cube_size: float = 0.05, + stack_xy_tolerance: float = 0.025, + stack_height_tolerance: float = 0.012, + target_xy_tolerance: float = 0.025, + settle_velocity_tolerance: float = 0.05, + ) -> None: + if render_mode not in {None, "human", "rgb_array"}: + raise ValueError(f"Unsupported render_mode: {render_mode}") + if top_cube_name not in self.CUBE_NAMES: + raise ValueError(f"Unknown top_cube_name: {top_cube_name!r}") + if base_cube_name not in self.CUBE_NAMES: + raise ValueError(f"Unknown base_cube_name: {base_cube_name!r}") + if top_cube_name == base_cube_name: + raise ValueError("top_cube_name and base_cube_name must differ.") + + super().__init__() + self.scene_file = scene_file + self.scene_path = resolve_scene_path(scene_file, version=version) + self.version = None if self.scene_path.parent == SCENES_ROOT else self.scene_path.parent.name + self.frame_skip = int(frame_skip) + self.render_mode = render_mode + self.workspace_bounds = self._validate_workspace_bounds(workspace_bounds) + self.min_cube_spacing = float(min_cube_spacing) + self.target_bounds = self._validate_workspace_bounds( + self.workspace_bounds if target_bounds is None else target_bounds + ) + self.min_target_cube_spacing = float(min_target_cube_spacing) + self.randomize_yaw = bool(randomize_yaw) + self.render_camera = render_camera + self.top_cube_name = top_cube_name + self.base_cube_name = base_cube_name + self.cube_size = float(cube_size) + self.stack_xy_tolerance = float(stack_xy_tolerance) + self.stack_height_tolerance = float(stack_height_tolerance) + self.target_xy_tolerance = float(target_xy_tolerance) + self.settle_velocity_tolerance = float(settle_velocity_tolerance) + + self.model = mujoco.MjModel.from_xml_path(str(self.scene_path)) + self.data = mujoco.MjData(self.model) + self._renderer: mujoco.Renderer | None = None + self._viewer: Any | None = None + + self._cube_qpos_adrs = { + cube_name: int(self.model.jnt_qposadr[self.model.joint(joint_name).id]) + for cube_name, joint_name in zip( + self.CUBE_NAMES, + self.CUBE_JOINT_NAMES, + strict=True, + ) + } + self._cube_qvel_adrs = { + cube_name: int(self.model.jnt_dofadr[self.model.joint(joint_name).id]) + for cube_name, joint_name in zip( + self.CUBE_NAMES, + self.CUBE_JOINT_NAMES, + strict=True, + ) + } + self._default_cube_qpos = { + cube_name: self.model.qpos0[qpos_adr : qpos_adr + 7].copy() + for cube_name, qpos_adr in self._cube_qpos_adrs.items() + } + self._target_body_id = mujoco.mj_name2id( + self.model, + mujoco.mjtObj.mjOBJ_BODY, + self.TARGET_BODY_NAME, + ) + if self._target_body_id < 0: + raise ValueError( + f"Scene {self.scene_path} is missing target body {self.TARGET_BODY_NAME!r}." + ) + self._target_mocap_id = int(self.model.body_mocapid[self._target_body_id]) + if self._target_mocap_id < 0: + raise ValueError( + f"Target body {self.TARGET_BODY_NAME!r} must be a mocap body." + ) + self._default_target_pos = self.model.body_pos[self._target_body_id].copy() + + self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(0,), dtype=np.float32) + obs = self._get_obs() + self.observation_space = spaces.Box( + low=-np.inf, + high=np.inf, + shape=obs.shape, + dtype=np.float64, + ) + + def _get_obs(self) -> np.ndarray: + return np.concatenate([self.data.qpos.copy(), self.data.qvel.copy()]) + + def _get_info(self) -> dict[str, Any]: + cube_pos, cube_quat, cube_qvel = self._cube_state() + success_info = self._stack_success_info(cube_pos, cube_qvel) + return { + "cube_pos": cube_pos, + "cube_quat": cube_quat, + "cube_qvel": cube_qvel, + "target_pos": self._target_pos().copy(), + "top_cube": self.top_cube_name, + "base_cube": self.base_cube_name, + **success_info, + } + + def reset( + self, + *, + seed: int | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[np.ndarray, dict[str, Any]]: + super().reset(seed=seed) + options = {} if options is None else dict(options) + mujoco.mj_resetData(self.model, self.data) + + if "qpos" in options: + qpos = np.asarray(options["qpos"], dtype=np.float64) + if qpos.shape != self.data.qpos.shape: + raise ValueError( + f"Expected qpos shape {self.data.qpos.shape}, got {qpos.shape}" + ) + self.data.qpos[:] = qpos + cube_positions = self._cube_positions_from_qpos(self.data.qpos) + else: + qpos = self.model.qpos0.copy() + cube_positions = self._reset_cubes_from_options(options, qpos=qpos) + self.data.qpos[:] = qpos + + self._reset_target_from_options(options, cube_positions) + + if "qvel" in options: + qvel = np.asarray(options["qvel"], dtype=np.float64) + if qvel.shape != self.data.qvel.shape: + raise ValueError( + f"Expected qvel shape {self.data.qvel.shape}, got {qvel.shape}" + ) + self.data.qvel[:] = qvel + else: + self.data.qvel[:] = 0.0 + + mujoco.mj_forward(self.model, self.data) + settle_steps = int(options.get("settle_steps", 0)) + if settle_steps > 0: + mujoco.mj_step(self.model, self.data, nstep=settle_steps) + + return self._get_obs(), self._get_info() + + def step( + self, + action: np.ndarray | None, + ) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]: + if action is not None: + resolved_action = np.asarray(action, dtype=np.float32) + if resolved_action.shape != self.action_space.shape: + raise ValueError( + f"Expected action shape {self.action_space.shape}, got {resolved_action.shape}" + ) + + mujoco.mj_step(self.model, self.data, nstep=self.frame_skip) + info = self._get_info() + return self._get_obs(), float(info["is_success"]), bool(info["is_success"]), False, info + + def render(self) -> np.ndarray | None: + if self.render_mode == "rgb_array": + if self._renderer is None: + self._renderer = mujoco.Renderer(self.model) + self._renderer.update_scene(self.data, camera=self.render_camera) + return self._renderer.render() + + if self.render_mode == "human": + if self._viewer is None: + from mujoco import viewer + + try: + self._viewer = viewer.launch_passive(self.model, self.data) + except RuntimeError as exc: + if sys.platform == "darwin" and "mjpython" in str(exc): + raise RuntimeError( + "On macOS, MuJoCo human rendering must be launched with " + "`mjpython`, not plain `python3`." + ) from exc + raise + mujoco.mjv_defaultFreeCamera(self.model, self._viewer.cam) + self._viewer.sync() + return None + + return None + + def close(self) -> None: + if self._renderer is not None: + self._renderer.close() + self._renderer = None + if self._viewer is not None: + self._viewer.close() + self._viewer = None + + def _cube_state( + self, + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, np.ndarray]]: + cube_pos = { + cube_name: self._cube_qpos(cube_name)[:3].copy() + for cube_name in self.CUBE_NAMES + } + cube_quat = { + cube_name: self._cube_qpos(cube_name)[3:7].copy() + for cube_name in self.CUBE_NAMES + } + cube_qvel = { + cube_name: self.data.qvel[qvel_adr : qvel_adr + 6].copy() + for cube_name, qvel_adr in self._cube_qvel_adrs.items() + } + return cube_pos, cube_quat, cube_qvel + + def _reset_cubes_from_options( + self, + options: Mapping[str, Any], + *, + qpos: np.ndarray | None = None, + ) -> dict[str, np.ndarray]: + cube_positions = options.get("cube_positions") + if cube_positions is None: + cube_positions = self._sample_cube_positions() + cube_quats = options.get("cube_quats", {}) + resolved_positions: dict[str, np.ndarray] = {} + + target_qpos = self.data.qpos if qpos is None else qpos + for cube_name in self.CUBE_NAMES: + cube_qpos = self._default_cube_qpos[cube_name].copy() + if cube_name in cube_positions: + cube_pos = np.asarray(cube_positions[cube_name], dtype=np.float64) + if cube_pos.shape != (3,): + raise ValueError( + f"Expected {cube_name} position shape (3,), got {cube_pos.shape}" + ) + cube_qpos[:3] = cube_pos + if cube_name in cube_quats: + cube_qpos[3:7] = self._normalize_quat( + np.asarray(cube_quats[cube_name], dtype=np.float64) + ) + elif self.randomize_yaw: + cube_qpos[3:7] = self._sample_yaw_quat() + + qpos_adr = self._cube_qpos_adrs[cube_name] + target_qpos[qpos_adr : qpos_adr + 7] = cube_qpos + resolved_positions[cube_name] = cube_qpos[:3].copy() + + return resolved_positions + + def _reset_target_from_options( + self, + options: Mapping[str, Any], + cube_positions: Mapping[str, np.ndarray], + ) -> None: + target_position = options.get("target_position") + if target_position is None: + target_position = self._sample_target_position(cube_positions) + target_pos = np.asarray(target_position, dtype=np.float64) + if target_pos.shape != (3,): + raise ValueError(f"Expected target_position shape (3,), got {target_pos.shape}") + self.data.mocap_pos[self._target_mocap_id] = target_pos + self.data.mocap_quat[self._target_mocap_id] = np.array( + [1.0, 0.0, 0.0, 0.0], + dtype=np.float64, + ) + + def _sample_cube_positions(self) -> dict[str, np.ndarray]: + sampled: list[np.ndarray] = [] + z_by_cube = { + cube_name: self._default_cube_qpos[cube_name][2] + for cube_name in self.CUBE_NAMES + } + for cube_name in self.CUBE_NAMES: + for _ in range(200): + xy = self.np_random.uniform( + low=self.workspace_bounds[:, 0], + high=self.workspace_bounds[:, 1], + ) + candidate = np.array([xy[0], xy[1], z_by_cube[cube_name]], dtype=np.float64) + if all( + np.linalg.norm(candidate[:2] - existing[:2]) >= self.min_cube_spacing + for existing in sampled + ): + sampled.append(candidate) + break + else: + raise RuntimeError( + "Unable to sample non-overlapping cube positions in the tabletop workspace." + ) + return dict(zip(self.CUBE_NAMES, sampled, strict=True)) + + def _sample_target_position( + self, + cube_positions: Mapping[str, np.ndarray], + ) -> np.ndarray: + cube_xy = [np.asarray(pos, dtype=np.float64)[:2] for pos in cube_positions.values()] + for _ in range(200): + xy = self.np_random.uniform( + low=self.target_bounds[:, 0], + high=self.target_bounds[:, 1], + ) + if all( + np.linalg.norm(xy - existing_xy) >= self.min_target_cube_spacing + for existing_xy in cube_xy + ): + return np.array( + [xy[0], xy[1], self._default_target_pos[2]], + dtype=np.float64, + ) + raise RuntimeError( + "Unable to sample a target position away from cubes in the tabletop workspace." + ) + + def _cube_positions_from_qpos(self, qpos: np.ndarray) -> dict[str, np.ndarray]: + return { + cube_name: qpos[qpos_adr : qpos_adr + 3].copy() + for cube_name, qpos_adr in self._cube_qpos_adrs.items() + } + + def _target_pos(self) -> np.ndarray: + return self.data.mocap_pos[self._target_mocap_id] + + def _stack_success_info( + self, + cube_pos: Mapping[str, np.ndarray], + cube_qvel: Mapping[str, np.ndarray], + ) -> dict[str, bool]: + top = cube_pos[self.top_cube_name] + base = cube_pos[self.base_cube_name] + top_vel = cube_qvel[self.top_cube_name][:3] + base_vel = cube_qvel[self.base_cube_name][:3] + target = self._target_pos() + + xy_close = np.linalg.norm(top[:2] - base[:2]) < self.stack_xy_tolerance + height_ok = abs((top[2] - base[2]) - self.cube_size) < self.stack_height_tolerance + target_xy_close = np.linalg.norm(base[:2] - target[:2]) < self.target_xy_tolerance + settled = ( + np.linalg.norm(top_vel) < self.settle_velocity_tolerance + and np.linalg.norm(base_vel) < self.settle_velocity_tolerance + ) + is_success = bool(xy_close and height_ok and target_xy_close and settled) + + return { + "xy_close": bool(xy_close), + "height_ok": bool(height_ok), + "target_xy_close": bool(target_xy_close), + "settled": bool(settled), + "is_success": is_success, + } + + def _sample_yaw_quat(self) -> np.ndarray: + yaw = float(self.np_random.uniform(low=-np.pi, high=np.pi)) + return np.array( + [np.cos(0.5 * yaw), 0.0, 0.0, np.sin(0.5 * yaw)], + dtype=np.float64, + ) + + def _cube_qpos(self, cube_name: str) -> np.ndarray: + qpos_adr = self._cube_qpos_adrs[cube_name] + return self.data.qpos[qpos_adr : qpos_adr + 7] + + @staticmethod + def _normalize_quat(quat: np.ndarray) -> np.ndarray: + if quat.shape != (4,): + raise ValueError(f"Expected quaternion shape (4,), got {quat.shape}") + norm = np.linalg.norm(quat) + if norm <= 0: + raise ValueError("Quaternion must have non-zero norm.") + return quat / norm + + @staticmethod + def _validate_workspace_bounds( + workspace_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]], + ) -> np.ndarray: + bounds = np.asarray(workspace_bounds, dtype=np.float64) + if bounds.shape != (2, 2): + raise ValueError(f"Expected workspace_bounds shape (2, 2), got {bounds.shape}") + if np.any(bounds[:, 0] >= bounds[:, 1]): + raise ValueError("Each workspace lower bound must be below its upper bound.") + return bounds + + +class OrcaArmCubeStacking(CubeStackingTabletop): + """OrcaArm cube stacking task built directly on the composed MuJoCo scene.""" + + DEFAULT_KEYFRAME = "orcaarm_home" + + def __init__( + self, + render_mode: str | None = None, + version: str | None = None, + *, + scene_file: str = "orcaarm_cube_stacking_cameras.xml", + frame_skip: int = 5, + actuator_names: Sequence[str] | None = None, + camera_names: Sequence[str] | None = default_orcaarm_camera_names(), + camera_width: int = 128, + camera_height: int = 128, + workspace_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] = CubeStackingTabletop.DEFAULT_WORKSPACE_BOUNDS, + min_cube_spacing: float = 0.09, + target_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] | None = None, + min_target_cube_spacing: float = 0.09, + randomize_yaw: bool = True, + render_camera: str = "topdown", + home_keyframe: str = DEFAULT_KEYFRAME, + top_cube_name: str = "red_cube", + base_cube_name: str = "blue_cube", + cube_size: float = 0.05, + stack_xy_tolerance: float = 0.025, + stack_height_tolerance: float = 0.012, + target_xy_tolerance: float = 0.025, + settle_velocity_tolerance: float = 0.05, + max_episode_steps: int = 200, + ) -> None: + self.home_keyframe = home_keyframe + self.max_episode_steps = int(max_episode_steps) + self.camera_names = tuple(camera_names or ()) + self.camera_width = int(camera_width) + self.camera_height = int(camera_height) + self._elapsed_steps = 0 + + super().__init__( + render_mode=render_mode, + version=version, + scene_file=scene_file, + frame_skip=frame_skip, + workspace_bounds=workspace_bounds, + min_cube_spacing=min_cube_spacing, + target_bounds=target_bounds, + min_target_cube_spacing=min_target_cube_spacing, + randomize_yaw=randomize_yaw, + render_camera=render_camera, + top_cube_name=top_cube_name, + base_cube_name=base_cube_name, + cube_size=cube_size, + stack_xy_tolerance=stack_xy_tolerance, + stack_height_tolerance=stack_height_tolerance, + target_xy_tolerance=target_xy_tolerance, + settle_velocity_tolerance=settle_velocity_tolerance, + ) + + self._home_keyframe_id = mujoco.mj_name2id( + self.model, + mujoco.mjtObj.mjOBJ_KEY, + self.home_keyframe, + ) + if self._home_keyframe_id < 0: + raise ValueError( + f"Scene {self.scene_path} is missing keyframe {self.home_keyframe!r}." + ) + + self.actuator_names = self._resolve_actuator_names(actuator_names) + self.actuator_ids = tuple(self.model.actuator(name).id for name in self.actuator_names) + action_low = self.model.actuator_ctrlrange[list(self.actuator_ids), 0].astype(np.float32) + action_high = self.model.actuator_ctrlrange[list(self.actuator_ids), 1].astype(np.float32) + self.action_space = spaces.Box(low=action_low, high=action_high, dtype=np.float32) + self._camera_renderer: mujoco.Renderer | None = None + + for camera_name in self.camera_names: + if mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name) < 0: + raise ValueError( + f"Scene {self.scene_path} is missing camera {camera_name!r}." + ) + + def _resolve_actuator_names( + self, + actuator_names: Sequence[str] | None, + ) -> tuple[str, ...]: + if actuator_names is None: + return tuple(self.model.actuator(actuator_id).name for actuator_id in range(self.model.nu)) + + resolved_names = tuple(actuator_names) + missing = [ + actuator_name + for actuator_name in resolved_names + if mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, actuator_name) < 0 + ] + if missing: + raise ValueError(f"Unknown actuator name(s): {missing}") + return resolved_names + + def reset( + self, + *, + seed: int | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[np.ndarray, dict[str, Any]]: + gym.Env.reset(self, seed=seed) + self._elapsed_steps = 0 + options = {} if options is None else dict(options) + + mujoco.mj_resetDataKeyframe(self.model, self.data, self._home_keyframe_id) + + if "qpos" in options: + qpos = np.asarray(options["qpos"], dtype=np.float64) + if qpos.shape != self.data.qpos.shape: + raise ValueError( + f"Expected qpos shape {self.data.qpos.shape}, got {qpos.shape}" + ) + self.data.qpos[:] = qpos + cube_positions = self._cube_positions_from_qpos(self.data.qpos) + else: + cube_positions = self._reset_cubes_from_options(options) + + self._reset_target_from_options(options, cube_positions) + + if "qvel" in options: + qvel = np.asarray(options["qvel"], dtype=np.float64) + if qvel.shape != self.data.qvel.shape: + raise ValueError( + f"Expected qvel shape {self.data.qvel.shape}, got {qvel.shape}" + ) + self.data.qvel[:] = qvel + else: + self.data.qvel[:] = 0.0 + cube_qvels = options.get("cube_qvels", {}) + for cube_name, cube_qvel in cube_qvels.items(): + qvel = np.asarray(cube_qvel, dtype=np.float64) + if qvel.shape != (6,): + raise ValueError( + f"Expected {cube_name} qvel shape (6,), got {qvel.shape}" + ) + qvel_adr = self._cube_qvel_adrs[cube_name] + self.data.qvel[qvel_adr : qvel_adr + 6] = qvel + + if "ctrl" in options: + ctrl = np.asarray(options["ctrl"], dtype=np.float64) + if ctrl.shape != self.data.ctrl.shape: + raise ValueError( + f"Expected ctrl shape {self.data.ctrl.shape}, got {ctrl.shape}" + ) + self.data.ctrl[:] = ctrl + + mujoco.mj_forward(self.model, self.data) + settle_steps = int(options.get("settle_steps", 0)) + if settle_steps > 0: + mujoco.mj_step(self.model, self.data, nstep=settle_steps) + + return self._get_obs(), self._get_info() + + def step( + self, + action: np.ndarray, + ) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]: + resolved_action = np.asarray(action, dtype=np.float32) + if resolved_action.shape != self.action_space.shape: + raise ValueError( + f"Expected action shape {self.action_space.shape}, got {resolved_action.shape}" + ) + clipped_action = np.clip(resolved_action, self.action_space.low, self.action_space.high) + self.data.ctrl[list(self.actuator_ids)] = clipped_action + + mujoco.mj_step(self.model, self.data, nstep=self.frame_skip) + self._elapsed_steps += 1 + + obs = self._get_obs() + info = self._get_info() + reward = float(info["is_success"]) + terminated = bool(info["is_success"]) + truncated = self._elapsed_steps >= self.max_episode_steps + return obs, reward, terminated, truncated, info + + def _get_info(self) -> dict[str, Any]: + info = super()._get_info() + info["elapsed_steps"] = self._elapsed_steps + return info + + def render_camera_observations(self) -> dict[str, np.ndarray]: + if not self.camera_names: + return {} + if self._camera_renderer is None: + self._camera_renderer = mujoco.Renderer( + self.model, + width=self.camera_width, + height=self.camera_height, + ) + images: dict[str, np.ndarray] = {} + for camera_name in self.camera_names: + self._camera_renderer.update_scene(self.data, camera=camera_name) + images[camera_name] = self._camera_renderer.render() + return images + + def close(self) -> None: + if self._camera_renderer is not None: + self._camera_renderer.close() + self._camera_renderer = None + super().close() + + +class OrcaPandaCubeStacking(OrcaArmCubeStacking): + """Single-arm OrcaPanda cube stacking task in the shared tabletop scene.""" + + DEFAULT_KEYFRAME = "orcapanda_home" + DEFAULT_CAMERA_NAMES = ( + "orcapanda_overview", + "topdown", + "angled", + "orcapanda_wrist_camera", + ) + + def __init__( + self, + render_mode: str | None = None, + version: str | None = None, + *, + scene_file: str = "orcapanda_cube_stacking.xml", + frame_skip: int = 5, + actuator_names: Sequence[str] | None = None, + camera_names: Sequence[str] | None = DEFAULT_CAMERA_NAMES, + camera_width: int = 128, + camera_height: int = 128, + workspace_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] = CubeStackingTabletop.DEFAULT_WORKSPACE_BOUNDS, + min_cube_spacing: float = 0.09, + target_bounds: np.ndarray | list[list[float]] | tuple[tuple[float, float], tuple[float, float]] | None = None, + min_target_cube_spacing: float = 0.09, + randomize_yaw: bool = True, + render_camera: str = "orcapanda_overview", + home_keyframe: str = DEFAULT_KEYFRAME, + top_cube_name: str = "red_cube", + base_cube_name: str = "blue_cube", + cube_size: float = 0.05, + stack_xy_tolerance: float = 0.025, + stack_height_tolerance: float = 0.012, + target_xy_tolerance: float = 0.025, + settle_velocity_tolerance: float = 0.05, + max_episode_steps: int = 200, + ) -> None: + super().__init__( + render_mode=render_mode, + version=version, + scene_file=scene_file, + frame_skip=frame_skip, + actuator_names=actuator_names, + camera_names=camera_names, + camera_width=camera_width, + camera_height=camera_height, + workspace_bounds=workspace_bounds, + min_cube_spacing=min_cube_spacing, + target_bounds=target_bounds, + min_target_cube_spacing=min_target_cube_spacing, + randomize_yaw=randomize_yaw, + render_camera=render_camera, + home_keyframe=home_keyframe, + top_cube_name=top_cube_name, + base_cube_name=base_cube_name, + cube_size=cube_size, + stack_xy_tolerance=stack_xy_tolerance, + stack_height_tolerance=stack_height_tolerance, + target_xy_tolerance=target_xy_tolerance, + settle_velocity_tolerance=settle_velocity_tolerance, + max_episode_steps=max_episode_steps, + ) class OrcaHandRightCubeOrientation(BaseOrcaHandEnv): @@ -171,14 +854,15 @@ def _resolve_initial_cube_quat(self, options: dict[str, Any]) -> np.ndarray: return self._default_cube_quat.copy() return self._sample_random_nonsolved_quaternion(self.np_random) - def _compose_ctrl_from_qpos(self) -> np.ndarray: - ctrl = np.zeros(self.model.nu, dtype=np.float32) - for actuator_id, qpos_idx in enumerate(self._actuator_qpos_indices): - ctrl[actuator_id] = float( + def _compose_ctrl_from_qpos(self, qpos: np.ndarray | None = None) -> np.ndarray: + source_qpos = self.data.qpos if qpos is None else np.asarray(qpos, dtype=np.float64) + ctrl = np.zeros(len(self.hand.config.joint_ids), dtype=np.float32) + for ctrl_idx, qpos_idx in enumerate(self.hand.config.actuator_qpos_indices): + ctrl[ctrl_idx] = float( np.clip( - self.data.qpos[qpos_idx], - self.action_low[actuator_id], - self.action_high[actuator_id], + source_qpos[qpos_idx], + self.action_low[ctrl_idx], + self.action_high[ctrl_idx], ) ) return ctrl @@ -261,7 +945,6 @@ def reset( options: dict[str, Any] | None = None, ) -> tuple[np.ndarray, dict[str, Any]]: gym.Env.reset(self, seed=seed) - mujoco.mj_resetData(self.model, self.data) self._elapsed_steps = 0 options = {} if options is None else dict(options) @@ -298,9 +981,10 @@ def reset( cube_quat = self._resolve_initial_cube_quat(options) - self.data.qpos[: self._cube_qpos_adr] = hand_qpos - self.data.qpos[self._cube_qpos_adr : self._cube_qpos_adr + 3] = cube_pos - self.data.qpos[self._cube_qpos_adr + 3 : self._cube_qpos_adr + 7] = cube_quat + qpos = self.model.qpos0.copy() + qpos[: self._cube_qpos_adr] = hand_qpos + qpos[self._cube_qpos_adr : self._cube_qpos_adr + 3] = cube_pos + qpos[self._cube_qpos_adr + 3 : self._cube_qpos_adr + 7] = cube_quat if full_qvel is not None: qvel = np.asarray(full_qvel, dtype=np.float64) @@ -308,43 +992,29 @@ def reset( raise ValueError( f"Expected qvel shape {self.data.qvel.shape}, got {qvel.shape}" ) - self.data.qvel[:] = qvel else: - self.data.qvel[:] = 0.0 + qvel = np.zeros_like(self.data.qvel) if "cube_qvel" in options: cube_qvel = np.asarray(options["cube_qvel"], dtype=np.float64) if cube_qvel.shape != (6,): raise ValueError( f"Expected cube_qvel shape (6,), got {cube_qvel.shape}" ) - self.data.qvel[self._cube_qvel_adr : self._cube_qvel_adr + 6] = cube_qvel + qvel[self._cube_qvel_adr : self._cube_qvel_adr + 6] = cube_qvel - self.data.ctrl[:] = self._compose_ctrl_from_qpos() - mujoco.mj_forward(self.model, self.data) + ctrl = self._compose_ctrl_from_qpos(qpos) + self.hand.reset(qpos=qpos, qvel=qvel, ctrl=ctrl) settle_steps = int(options.get("settle_steps", 0)) for _ in range(settle_steps): - mujoco.mj_step(self.model, self.data) - self.data.ctrl[:] = self._compose_ctrl_from_qpos() - - mujoco.mj_forward(self.model, self.data) - - if self.render_mode == "human": - self.render() + self.hand.step(nstep=1) return self._get_obs(), self._get_info() def step( self, action: np.ndarray ) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]: - action = np.asarray(action, dtype=np.float32) - if action.shape != self.action_space.shape: - raise ValueError( - f"Expected action shape {self.action_space.shape}, got {action.shape}" - ) - - self.data.ctrl[:] = np.clip(action, self.action_low, self.action_high) - mujoco.mj_step(self.model, self.data, nstep=self.frame_skip) + self.hand.step(action) self._elapsed_steps += 1 obs = self._get_obs() @@ -353,9 +1023,6 @@ def step( truncated = self._get_truncated() info = self._get_info() - if self.render_mode == "human": - self.render() - return obs, reward, terminated, truncated, info @staticmethod diff --git a/src/orca_sim/versions.py b/src/orca_sim/versions.py index fb5bfe6..4b2ac6c 100644 --- a/src/orca_sim/versions.py +++ b/src/orca_sim/versions.py @@ -36,15 +36,21 @@ def resolve_version(version: str | None = None) -> str: def resolve_scene_path(scene_file: str, *, version: str | None = None) -> Path: + unversioned_scene_path = SCENES_ROOT / scene_file if version is not None: resolved_version = resolve_version(version) scene_path = SCENES_ROOT / resolved_version / scene_file - if not scene_path.exists(): - raise FileNotFoundError( - f"Embodiment version '{resolved_version}' is missing required scene file: " - f"{scene_file}" - ) - return scene_path + if scene_path.exists(): + return scene_path + if unversioned_scene_path.exists(): + return unversioned_scene_path + raise FileNotFoundError( + f"Embodiment version '{resolved_version}' is missing required scene file " + f"or unversioned scene: {scene_file}" + ) + + if unversioned_scene_path.exists(): + return unversioned_scene_path for candidate_version in _scene_resolution_order(): scene_path = SCENES_ROOT / candidate_version / scene_file @@ -53,7 +59,7 @@ def resolve_scene_path(scene_file: str, *, version: str | None = None) -> Path: known_versions = ", ".join(list_versions()) or "none" raise FileNotFoundError( - f"No embodiment version provides scene file '{scene_file}'. " + f"No unversioned scene or embodiment version provides scene file '{scene_file}'. " f"Available versions: {known_versions}" ) diff --git a/tests/test_envs.py b/tests/test_envs.py index d7dec72..e68a9a3 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -1,7 +1,9 @@ import numpy as np import pytest -from orca_sim import OrcaHandCombined, OrcaHandLeft, OrcaHandRight +from orca_core.base_hand import BaseHand + +from orca_sim import OrcaHandCombined, OrcaHandLeft, OrcaHandRight, SimOrcaHand @pytest.mark.parametrize( @@ -39,6 +41,18 @@ def test_env_reset_and_step_smoke( env.close() +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_envs_compose_shared_sim_hand_backend(version: str) -> None: + env = OrcaHandRight(version=version) + try: + assert isinstance(env.hand, SimOrcaHand) + assert isinstance(env.hand, BaseHand) + assert env.hand.version == version + assert env.hand.scene_path == env.scene_path + finally: + env.close() + + def test_reset_accepts_explicit_qpos_and_qvel() -> None: env = OrcaHandRight() try: @@ -55,12 +69,62 @@ def test_reset_accepts_explicit_qpos_and_qvel() -> None: env.close() +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_sim_hand_clamps_joint_commands(version: str) -> None: + hand = SimOrcaHand(scene_file="scene_right.xml", version=version) + try: + hand.reset() + joint_name = hand.config.joint_ids[0] + hand.set_joint_positions({joint_name: hand.action_high[0] + 10.0}) + joint_positions = hand.get_joint_position().as_dict() + + assert joint_positions[joint_name] == pytest.approx(hand.action_high[0]) + finally: + hand.close() + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_sim_hand_preserves_unspecified_joint_commands(version: str) -> None: + hand = SimOrcaHand(scene_file="scene_right.xml", version=version) + try: + hand.reset() + first_joint, second_joint = hand.config.joint_ids[:2] + + hand.set_joint_positions( + { + first_joint: float(hand.action_low[0]), + second_joint: float(hand.action_high[1]), + } + ) + hand.set_joint_positions({second_joint: 0.0}) + + joint_positions = hand.get_joint_position().as_dict() + assert joint_positions[first_joint] == pytest.approx(hand.action_low[0]) + assert joint_positions[second_joint] == pytest.approx(0.0) + finally: + hand.close() + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_sim_hand_steps_simulation_without_gym(version: str) -> None: + hand = SimOrcaHand(scene_file="scene_right.xml", version=version) + try: + obs = hand.reset() + next_obs = hand.step(hand.action_high + 10.0) # overwrap all motors + + assert obs.shape == next_obs.shape == (hand.data.qpos.size + hand.data.qvel.size,) + np.testing.assert_allclose(hand.data.ctrl[list(hand.config.actuator_ids)], hand.action_high) + finally: + hand.close() + + def test_step_clips_actions_to_actuator_limits() -> None: env = OrcaHandLeft() try: env.reset() env.step(env.action_high + 10.0) - np.testing.assert_allclose(env.data.ctrl, env.action_high) + # indexing of joint ids is specified by the hand config + np.testing.assert_allclose(env.data.ctrl[list(env.hand.config.actuator_ids)], env.action_high) finally: env.close() diff --git a/tests/test_registry.py b/tests/test_registry.py index 042cb52..4b29aa3 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -22,7 +22,13 @@ def test_register_envs_is_idempotent_and_envs_can_be_made() -> None: ("OrcaHandRightCubeOrientation-v2", (51,), (17,), False), ("OrcaHandCombined-v1", (68,), (34,), True), ("OrcaHandCombined-v2", (68,), (34,), True), + ("CubeStackingTabletop", (26,), (0,), False), + ("OrcaArmCubeStacking", (114,), (44,), False), + ("OrcaPandaCubeStacking", (74,), (24,), False), ]: + if env_id in {"OrcaArmCubeStacking", "OrcaPandaCubeStacking"}: + pytest.importorskip("orca_arm") + assert env_id in gym.registry env = gym.make(env_id) @@ -33,6 +39,12 @@ def test_register_envs_is_idempotent_and_envs_can_be_made() -> None: assert env.action_space.shape == action_shape if expect_empty_info: assert info == {} + elif env_id in { + "CubeStackingTabletop", + "OrcaArmCubeStacking", + "OrcaPandaCubeStacking", + }: + assert set(info["cube_pos"]) == {"red_cube", "blue_cube"} else: assert "red_face_up_alignment" in info finally: diff --git a/tests/test_sim_hand.py b/tests/test_sim_hand.py new file mode 100644 index 0000000..eb5d4f3 --- /dev/null +++ b/tests/test_sim_hand.py @@ -0,0 +1,504 @@ +"""Integration tests for the runtime hand API stack. + +These tests validate concordance between: +- MuJoCo scene assets +- SimOrcaHandConfig metadata derivation +- SimOrcaHand's shared BaseHand contract +- Gym env wrappers built on top of SimOrcaHand +- Task env wrappers built on top of SimOrcaHand +""" + +import mujoco +import numpy as np +import pytest +from orca_core.joint_position import OrcaJointPositions + +from orca_sim import ( + CubeStackingTabletop, + OrcaArmCubeStacking, + OrcaHandRight, + OrcaHandRightCubeOrientation, + OrcaPandaCubeStacking, + SimOrcaHand, + SimOrcaHandConfig, +) +from orca_sim.joint_mapping import ( + canonical_single_hand_joint_ids, + default_joint_name_to_scene_joint_name, +) +from orca_sim.versions import resolve_scene_path + + +def _qpos_indices_from_model(model: mujoco.MjModel) -> list[int]: + return [ + int(model.jnt_qposadr[int(model.actuator_trnid[actuator_id, 0])]) + for actuator_id in range(model.nu) + ] + + +def _build_valid_qpos(model: mujoco.MjModel) -> np.ndarray: + qpos = model.qpos0.copy() + ctrl_mid = model.actuator_ctrlrange.mean(axis=1) + for qpos_idx, value in zip(_qpos_indices_from_model(model), ctrl_mid, strict=True): + qpos[qpos_idx] = value + return qpos + + +def _raw_reset( + scene_path, + config: SimOrcaHandConfig, + *, + qpos: np.ndarray | None = None, + qvel: np.ndarray | None = None, + ctrl: np.ndarray | None = None, +) -> tuple[mujoco.MjModel, mujoco.MjData, np.ndarray]: + model = mujoco.MjModel.from_xml_path(str(scene_path)) + data = mujoco.MjData(model) + mujoco.mj_resetData(model, data) + + if qpos is not None: + data.qpos[:] = np.asarray(qpos, dtype=np.float64) + if qvel is not None: + data.qvel[:] = np.asarray(qvel, dtype=np.float64) + + if ctrl is None: + ctrl = np.array( + [data.qpos[qpos_idx] for qpos_idx in config.actuator_qpos_indices], + dtype=np.float32, + ) + else: + ctrl = np.asarray(ctrl, dtype=np.float32) + ctrl = np.clip(ctrl, np.asarray(config.action_low), np.asarray(config.action_high)) + + # Set actuators to control values + for actuator_id, value in zip(config.actuator_ids, ctrl, strict=True): + data.ctrl[actuator_id] = float(value) + + mujoco.mj_forward(model, data) + obs = np.concatenate([data.qpos.copy(), data.qvel.copy()]) + + return model, data, obs + + +@pytest.mark.parametrize( + ("scene_file", "version", "expected_type"), + [ + ("scene_left.xml", "v1", "left"), + ("scene_left.xml", "v2", "left"), + ("scene_right.xml", "v1", "right"), + ("scene_right.xml", "v2", "right"), + ("scene_combined.xml", "v1", None), + ("scene_combined.xml", "v2", None), + ], +) +def test_sim_hand_config_matches_scene_metadata( + scene_file: str, version: str, expected_type: str | None +) -> None: + """Tests mujoco models follows specs from the higher-level hand config. In particular, it ensures: + - Correct hand type is inferred from joint names. + - Correct joint name to scene joint name mapping. + - Correct derivation of: + - joint ids, + - scene joint names, + - actuator ids, actuator qpos indices, actuator qvel indices, + - action bounds [low, high] + - neutral position + """ + + """Mujoco model""" + scene_path = resolve_scene_path(scene_file, version=version) + model = mujoco.MjModel.from_xml_path(str(scene_path)) + + """SimOrcaHandConfig""" + config = SimOrcaHandConfig.from_config_path(scene_file=scene_file, version=version) + + resolved_type, expected_mapping = default_joint_name_to_scene_joint_name( + scene_file=scene_file, + version=version, + ) + expected_scene_joint_names = list(expected_mapping.values()) + + scene_joint_to_actuator_id = {} + scene_joint_to_qpos_idx = {} + scene_joint_to_qvel_idx = {} + for actuator_id in range(model.nu): + joint_id = int(model.actuator_trnid[actuator_id, 0]) + joint_name = model.joint(joint_id).name + scene_joint_to_actuator_id[joint_name] = actuator_id + scene_joint_to_qpos_idx[joint_name] = int(model.jnt_qposadr[joint_id]) + scene_joint_to_qvel_idx[joint_name] = int(model.jnt_dofadr[joint_id]) + + assert config.scene_path == str(scene_path) + assert config.scene_file == scene_file + assert config.version == version + assert resolved_type == expected_type + assert config.type == expected_type + assert list(config.scene_joint_names) == expected_scene_joint_names + assert list(config.actuator_ids) == [ + scene_joint_to_actuator_id[joint_name] for joint_name in expected_scene_joint_names + ] + assert list(config.actuator_qpos_indices) == [ + scene_joint_to_qpos_idx[joint_name] for joint_name in expected_scene_joint_names + ] + assert list(config.actuator_qvel_indices) == [ + scene_joint_to_qvel_idx[joint_name] for joint_name in expected_scene_joint_names + ] + assert np.allclose( + np.asarray(config.action_low), + [model.actuator_ctrlrange[actuator_id, 0] for actuator_id in config.actuator_ids], + ) + assert np.allclose( + np.asarray(config.action_high), + [model.actuator_ctrlrange[actuator_id, 1] for actuator_id in config.actuator_ids], + ) + + for joint_name, scene_joint_name in expected_mapping.items(): + qpos_idx = scene_joint_to_qpos_idx[scene_joint_name] + assert config.neutral_position[joint_name] == pytest.approx(model.qpos0[qpos_idx]) + + +@pytest.mark.parametrize( + ("scene_file", "version", "hand_type"), + [ + ("scene_right.xml", "v1", "right"), + ("scene_right.xml", "v2", "right"), + ("scene_left.xml", "v1", "left"), + ("scene_left.xml", "v2", "left"), + ], +) +def test_sim_hand_uses_orca_core_canonical_order_single_hand( + scene_file: str, version: str, hand_type: str +) -> None: + expected = list(canonical_single_hand_joint_ids(version=version, hand_type=hand_type)) + config = SimOrcaHandConfig.from_config_path(scene_file=scene_file, version=version) + assert config.joint_ids == expected + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_sim_hand_uses_orca_core_canonical_order_combined(version: str) -> None: + canonical = list(canonical_single_hand_joint_ids(version=version, hand_type="right")) + config = SimOrcaHandConfig.from_config_path(scene_file="scene_combined.xml", version=version) + assert config.joint_ids == [ + *[f"left_{joint}" for joint in canonical], + *[f"right_{joint}" for joint in canonical], + ] + + +@pytest.mark.parametrize( + ("scene_file", "version"), + [ + ("scene_right.xml", "v1"), + ("scene_right.xml", "v2"), + ("scene_right_cube_orientation.xml", "v1"), + ("scene_right_cube_orientation.xml", "v2"), + ], +) +def test_sim_hand_reset_matches_raw_mujoco(scene_file: str, version: str) -> None: + config = SimOrcaHandConfig.from_config_path(scene_file=scene_file, version=version) + scene_path = resolve_scene_path(scene_file, version=version) + model = mujoco.MjModel.from_xml_path(str(scene_path)) + qpos = _build_valid_qpos(model) + qvel = np.linspace(-0.2, 0.2, model.nv, dtype=np.float64) + ctrl = np.linspace(-10.0, 10.0, len(config.joint_ids), dtype=np.float32) + + _, raw_data, raw_obs = _raw_reset(scene_path, config, qpos=qpos, qvel=qvel, ctrl=ctrl) + + hand = SimOrcaHand(scene_file=scene_file, version=version) + try: + obs = hand.reset(qpos=qpos, qvel=qvel, ctrl=ctrl) + + np.testing.assert_allclose(obs, raw_obs) + np.testing.assert_allclose(hand.data.qpos, raw_data.qpos) + np.testing.assert_allclose(hand.data.qvel, raw_data.qvel) + np.testing.assert_allclose(hand.data.ctrl, raw_data.ctrl) + finally: + hand.close() + + +@pytest.mark.parametrize( + ("scene_file", "version"), + [ + ("scene_right.xml", "v1"), + ("scene_right.xml", "v2"), + ("scene_right_cube_orientation.xml", "v1"), + ("scene_right_cube_orientation.xml", "v2"), + ], +) +def test_sim_hand_step_matches_raw_mujoco(scene_file: str, version: str) -> None: + config = SimOrcaHandConfig.from_config_path(scene_file=scene_file, version=version) + scene_path = resolve_scene_path(scene_file, version=version) + model = mujoco.MjModel.from_xml_path(str(scene_path)) + qpos = _build_valid_qpos(model) + qvel = np.linspace(-0.1, 0.1, model.nv, dtype=np.float64) + action = np.linspace(-5.0, 5.0, len(config.joint_ids), dtype=np.float32) + + raw_model, raw_data, _ = _raw_reset(scene_path, config, qpos=qpos, qvel=qvel) + clipped_action = np.clip(action, np.asarray(config.action_low), np.asarray(config.action_high)) + for actuator_id, value in zip(config.actuator_ids, clipped_action, strict=True): + raw_data.ctrl[actuator_id] = float(value) + mujoco.mj_step(raw_model, raw_data, nstep=5) + raw_obs = np.concatenate([raw_data.qpos.copy(), raw_data.qvel.copy()]) + + hand = SimOrcaHand(scene_file=scene_file, version=version) + try: + hand.reset(qpos=qpos, qvel=qvel) + obs = hand.step(action) + + np.testing.assert_allclose(obs, raw_obs) + np.testing.assert_allclose(hand.data.qpos, raw_data.qpos) + np.testing.assert_allclose(hand.data.qvel, raw_data.qvel) + np.testing.assert_allclose(hand.data.ctrl, raw_data.ctrl) + finally: + hand.close() + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_shared_base_hand_helpers_work_for_sim_hand(version: str) -> None: + hand = SimOrcaHand(scene_file="scene_right.xml", version=version) + try: + hand.reset() + + target = np.linspace(-0.05, 0.05, len(hand.config.joint_ids), dtype=np.float64) + hand.set_joint_positions(target) + typed_joint_positions = hand.get_joint_position() + + assert isinstance(typed_joint_positions, OrcaJointPositions) + np.testing.assert_allclose( + typed_joint_positions.as_array(hand.config.joint_ids), + target, + ) + + hand.set_zero_position() + np.testing.assert_allclose( + hand.get_joint_position().as_array(hand.config.joint_ids), + np.zeros(len(hand.config.joint_ids), dtype=np.float64), + ) + + hand.set_neutral_position() + np.testing.assert_allclose( + hand.get_joint_position().as_array(hand.config.joint_ids), + np.array( + [hand.config.neutral_position[joint] for joint in hand.config.joint_ids], + dtype=np.float64, + ), + ) + finally: + hand.close() + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_base_env_delegates_reset_and_step_to_sim_hand(mocker, version: str) -> None: + env = OrcaHandRight(version=version) + try: + hand_reset = mocker.patch.object(env.hand, "reset", wraps=env.hand.reset) + env.reset() + hand_reset.assert_called_once() + mocker.stop(hand_reset) + + hand_step = mocker.patch.object(env.hand, "step", wraps=env.hand.step) + env.step(env.action_space.sample()) + hand_step.assert_called_once() + finally: + env.close() + + +@pytest.mark.parametrize("version", ["v1", "v2"]) +def test_task_env_delegates_reset_and_step_to_sim_hand(mocker, version: str) -> None: + env = OrcaHandRightCubeOrientation(version=version) + try: + hand_reset = mocker.patch.object(env.hand, "reset", wraps=env.hand.reset) + hand_step = mocker.patch.object(env.hand, "step", wraps=env.hand.step) + env.reset(options={"settle_steps": 2}) + hand_reset.assert_called_once() + assert hand_step.call_count == 2 + mocker.stop(hand_reset) + mocker.stop(hand_step) + + hand_step = mocker.patch.object(env.hand, "step", wraps=env.hand.step) + env.step(env.action_space.sample()) + hand_step.assert_called_once() + finally: + env.close() + + +def test_cube_stacking_tabletop_loads_and_randomizes_cubes() -> None: + env = CubeStackingTabletop() + try: + obs_0, info_0 = env.reset(seed=0) + obs_1, info_1 = env.reset(seed=1) + + assert env.version is None + assert obs_0.shape == env.observation_space.shape + assert obs_1.shape == env.observation_space.shape + assert env.action_space.shape == (0,) + assert set(info_0["cube_pos"]) == {"red_cube", "blue_cube"} + assert info_0["target_pos"].shape == (3,) + + red_0 = info_0["cube_pos"]["red_cube"] + blue_0 = info_0["cube_pos"]["blue_cube"] + red_1 = info_1["cube_pos"]["red_cube"] + blue_1 = info_1["cube_pos"]["blue_cube"] + assert np.linalg.norm(red_0[:2] - blue_0[:2]) >= env.min_cube_spacing + assert np.linalg.norm(red_1[:2] - blue_1[:2]) >= env.min_cube_spacing + assert not np.allclose(red_0[:2], red_1[:2]) + assert not np.allclose(blue_0[:2], blue_1[:2]) + assert not np.allclose(info_0["target_pos"][:2], info_1["target_pos"][:2]) + + env.step(env.action_space.sample()) + finally: + env.close() + + +def test_orcaarm_cube_stacking_uses_home_keyframe_and_all_actuators() -> None: + pytest.importorskip("orca_arm") + env = OrcaArmCubeStacking() + try: + obs_0, info_0 = env.reset(seed=0) + obs_1, info_1 = env.reset(seed=1) + + assert obs_0.shape == env.observation_space.shape + assert obs_1.shape == env.observation_space.shape + assert env.action_space.shape == (env.model.nu,) + assert len(env.actuator_names) == env.model.nu + assert set(info_0["cube_pos"]) == {"red_cube", "blue_cube"} + assert not info_0["is_success"] + assert not np.allclose( + info_0["cube_pos"]["red_cube"][:2], + info_1["cube_pos"]["red_cube"][:2], + ) + + obs, reward, terminated, truncated, info = env.step(env.action_space.sample()) + assert obs.shape == env.observation_space.shape + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + assert "is_success" in info + finally: + env.close() + + +def test_orcaarm_cube_stacking_renders_camera_observations() -> None: + pytest.importorskip("orca_arm") + env = OrcaArmCubeStacking(camera_width=64, camera_height=48) + try: + env.reset(seed=0) + images = env.render_camera_observations() + + assert set(images) == { + "chest_table_camera", + "left_wrist_camera", + "right_wrist_camera", + } + for image in images.values(): + assert image.shape == (48, 64, 3) + assert image.dtype == np.uint8 + finally: + env.close() + + +def test_orcaarm_cube_stacking_supports_ordered_actuator_subset() -> None: + pytest.importorskip("orca_arm") + actuator_names = ("act_openarm_right_joint1", "act_openarm_left_joint1") + env = OrcaArmCubeStacking(actuator_names=actuator_names) + try: + env.reset(seed=0) + + assert env.actuator_names == actuator_names + assert env.action_space.shape == (2,) + action = np.array([0.25, -0.5], dtype=np.float32) + env.step(action) + + expected_ids = [env.model.actuator(name).id for name in actuator_names] + np.testing.assert_allclose(env.data.ctrl[expected_ids], action) + finally: + env.close() + + +def test_orcaarm_cube_stacking_success_predicate() -> None: + pytest.importorskip("orca_arm") + env = OrcaArmCubeStacking(randomize_yaw=False) + try: + _, info = env.reset( + options={ + "cube_positions": { + "blue_cube": np.array([0.45, 0.0, 0.35], dtype=np.float64), + "red_cube": np.array([0.45, 0.0, 0.40], dtype=np.float64), + }, + "target_position": np.array([0.45, 0.0, 0.326], dtype=np.float64), + } + ) + + assert info["xy_close"] + assert info["height_ok"] + assert info["target_xy_close"] + assert info["settled"] + assert info["is_success"] + finally: + env.close() + + +def test_orcapanda_cube_stacking_uses_home_keyframe_and_all_actuators() -> None: + pytest.importorskip("orca_arm") + env = OrcaPandaCubeStacking() + try: + obs_0, info_0 = env.reset(seed=0) + obs_1, info_1 = env.reset(seed=1) + + assert obs_0.shape == env.observation_space.shape + assert obs_1.shape == env.observation_space.shape + assert env.action_space.shape == (env.model.nu,) + assert env.model.nu == 24 + assert len(env.actuator_names) == env.model.nu + assert set(info_0["cube_pos"]) == {"red_cube", "blue_cube"} + assert not info_0["is_success"] + assert not np.allclose( + info_0["cube_pos"]["red_cube"][:2], + info_1["cube_pos"]["red_cube"][:2], + ) + + obs, reward, terminated, truncated, info = env.step(env.action_space.sample()) + assert obs.shape == env.observation_space.shape + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + assert "is_success" in info + finally: + env.close() + + +def test_orcapanda_cube_stacking_renders_camera_observations() -> None: + pytest.importorskip("orca_arm") + env = OrcaPandaCubeStacking(camera_width=64, camera_height=48) + try: + env.reset(seed=0) + images = env.render_camera_observations() + + assert set(images) == { + "orcapanda_overview", + "topdown", + "angled", + "orcapanda_wrist_camera", + } + for image in images.values(): + assert image.shape == (48, 64, 3) + assert image.dtype == np.uint8 + finally: + env.close() + + +def test_orcapanda_cube_stacking_supports_ordered_actuator_subset() -> None: + pytest.importorskip("orca_arm") + actuator_names = ("act_panda_joint1", "act_panda_joint2") + env = OrcaPandaCubeStacking(actuator_names=actuator_names) + try: + env.reset(seed=0) + + assert env.actuator_names == actuator_names + assert env.action_space.shape == (2,) + action = np.array([0.25, -0.5], dtype=np.float32) + env.step(action) + + expected_ids = [env.model.actuator(name).id for name in actuator_names] + np.testing.assert_allclose(env.data.ctrl[expected_ids], action) + finally: + env.close() diff --git a/tests/test_versions.py b/tests/test_versions.py index e84842a..a37cd7f 100644 --- a/tests/test_versions.py +++ b/tests/test_versions.py @@ -1,6 +1,12 @@ +import os +import xml.etree.ElementTree as ET from pathlib import Path +import shutil +import subprocess +import sys import mujoco +import numpy as np import pytest from orca_sim.versions import ( @@ -31,6 +37,10 @@ def test_version_discovery_defaults_to_latest() -> None: "scene_left_extended.xml", "scene_right_extended.xml", "scene_combined_extended.xml", + "cube_stacking.xml", + "orcaarm_cube_stacking.xml", + "orcaarm_cube_stacking_cameras.xml", + "orcapanda_cube_stacking.xml", ], ) def test_scene_paths_exist_for_each_scene(scene_file: str) -> None: @@ -46,6 +56,163 @@ def test_resolve_version_rejects_unknown_versions() -> None: resolve_version("does-not-exist") +def test_unversioned_scene_resolves_from_scenes_root() -> None: + scene_path = resolve_scene_path("cube_stacking.xml") + + assert scene_path == PACKAGE_ROOT / "scenes" / "cube_stacking.xml" + + +def test_unversioned_scene_can_resolve_with_version_hint() -> None: + scene_path = resolve_scene_path("cube_stacking.xml", version="v2") + + assert scene_path == PACKAGE_ROOT / "scenes" / "cube_stacking.xml" + + +def test_orcaarm_cube_stacking_scene_references_orca_arm_mjcf_path() -> None: + orca_arm = pytest.importorskip("orca_arm") + scene_path = PACKAGE_ROOT / "scenes" / "orcaarm_cube_stacking.xml" + include_file = ET.parse(scene_path).getroot().findall("include")[-1].get("file") + + assert (scene_path.parent / include_file).resolve() == Path(orca_arm.MJCF_PATH).resolve() + + +def test_orcaarm_cube_stacking_scene_loads_directly() -> None: + pytest.importorskip("orca_arm") + + mujoco.MjModel.from_xml_path( + str((PACKAGE_ROOT / "scenes" / "orcaarm_cube_stacking.xml").resolve()) + ) + + +def test_orcaarm_cube_stacking_camera_scene_loads_directly() -> None: + pytest.importorskip("orca_arm") + + model = mujoco.MjModel.from_xml_path( + str((PACKAGE_ROOT / "scenes" / "orcaarm_cube_stacking_cameras.xml").resolve()) + ) + camera_names = {model.camera(camera_id).name for camera_id in range(model.ncam)} + assert { + "chest_table_camera", + "left_wrist_camera", + "right_wrist_camera", + }.issubset(camera_names) + + +def test_orcapanda_cube_stacking_scene_loads_directly() -> None: + pytest.importorskip("orca_arm") + + model = mujoco.MjModel.from_xml_path( + str((PACKAGE_ROOT / "scenes" / "orcapanda_cube_stacking.xml").resolve()) + ) + + assert model.nu == 24 + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "panda_link0") >= 0 + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "orcapanda_home") >= 0 + assert ( + mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, "orcapanda_wrist_camera") + >= 0 + ) + + +def test_orcapanda_cube_stacking_home_keyframe_sets_ready_pose() -> None: + pytest.importorskip("orca_arm") + q_home = np.array( + [-0.1, -1.6, -0.1, -3.0718, -0.15, 2.85, -1.4027], + dtype=np.float64, + ) + model = mujoco.MjModel.from_xml_path( + str((PACKAGE_ROOT / "scenes" / "orcapanda_cube_stacking.xml").resolve()) + ) + data = mujoco.MjData(model) + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "orcapanda_home") + + mujoco.mj_resetDataKeyframe(model, data, key_id) + + arm_joint_names = [f"panda_joint{i}" for i in range(1, 8)] + arm_qpos = np.array( + [ + data.qpos[int(model.jnt_qposadr[model.joint(joint_name).id])] + for joint_name in arm_joint_names + ], + dtype=np.float64, + ) + arm_ctrl = [] + for joint_name in arm_joint_names: + for actuator_id in range(model.nu): + joint_id = int(model.actuator_trnid[actuator_id, 0]) + if model.joint(joint_id).name == joint_name: + arm_ctrl.append(data.ctrl[actuator_id]) + break + + np.testing.assert_allclose(arm_qpos, q_home) + np.testing.assert_allclose(np.asarray(arm_ctrl, dtype=np.float64), q_home) + + +def test_orcaarm_cube_stacking_home_keyframe_sets_arm_pose() -> None: + pytest.importorskip("orca_arm") + q_home = np.array( + [ + 0.0, + 0.004, + 0.0, + 1.520, + 1.570796, + 0.0, + 0.005, + 0.0, + 1.530, + -1.570796, + ], + dtype=np.float64, + ) + model = mujoco.MjModel.from_xml_path( + str((PACKAGE_ROOT / "scenes" / "orcaarm_cube_stacking.xml").resolve()) + ) + data = mujoco.MjData(model) + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "orcaarm_home") + + mujoco.mj_resetDataKeyframe(model, data, key_id) + + arm_joint_names = [ + *[f"openarm_left_joint{i}" for i in range(1, 6)], + *[f"openarm_right_joint{i}" for i in range(1, 6)], + ] + arm_qpos = np.array( + [ + data.qpos[int(model.jnt_qposadr[model.joint(joint_name).id])] + for joint_name in arm_joint_names + ], + dtype=np.float64, + ) + arm_ctrl = [] + arm_joint_ids = set() + arm_actuator_ids = set() + for joint_name in arm_joint_names: + arm_joint_ids.add(model.joint(joint_name).id) + for actuator_id in range(model.nu): + joint_id = int(model.actuator_trnid[actuator_id, 0]) + if model.joint(joint_id).name == joint_name: + arm_ctrl.append(data.ctrl[actuator_id]) + arm_actuator_ids.add(actuator_id) + break + np.testing.assert_allclose(arm_qpos, q_home) + np.testing.assert_allclose(np.asarray(arm_ctrl, dtype=np.float64), q_home) + + for joint_id in range(model.njnt): + if joint_id in arm_joint_ids: + continue + joint_type = int(model.jnt_type[joint_id]) + qpos_width = 7 if joint_type == mujoco.mjtJoint.mjJNT_FREE else 1 + qpos_adr = int(model.jnt_qposadr[joint_id]) + np.testing.assert_allclose( + data.qpos[qpos_adr : qpos_adr + qpos_width], + model.qpos0[qpos_adr : qpos_adr + qpos_width], + ) + for actuator_id in range(model.nu): + if actuator_id not in arm_actuator_ids: + assert data.ctrl[actuator_id] == pytest.approx(0.0) + + @pytest.mark.parametrize( "scene_path", [ @@ -60,7 +227,91 @@ def test_resolve_version_rejects_unknown_versions() -> None: "src/orca_sim/scenes/v2/scene_right.xml", "src/orca_sim/scenes/v2/scene_combined.xml", "src/orca_sim/scenes/v2/scene_right_cube_orientation.xml", + "src/orca_sim/scenes/cube_stacking.xml", ], ) def test_scene_paths_load_directly(scene_path: str) -> None: mujoco.MjModel.from_xml_path(str(Path(scene_path).resolve())) + + +@pytest.mark.slow +@pytest.mark.skipif( + not os.environ.get("CI"), + reason="Packaging smoke test is slow; run it in CI or locally with CI=1.", +) +def test_built_wheel_installs_and_resets_right_hand(tmp_path: Path) -> None: + root = Path(__file__).resolve().parents[1] + project_dir = tmp_path / "project" + wheel_dir = tmp_path / "wheelhouse" + shutil.copytree( + root, + project_dir, + ignore=shutil.ignore_patterns( + ".git", + ".pytest_cache", + "__pycache__", + "*.pyc", + "build", + "dist", + "*.egg-info", + ), + ) + wheel_dir.mkdir() + + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "wheel", + ".", + "--no-deps", + "--no-build-isolation", + "--wheel-dir", + str(wheel_dir), + ], + cwd=project_dir, + check=True, + capture_output=True, + text=True, + ) + + wheel_path = next(wheel_dir.glob("orca_sim-*.whl")) + venv_dir = tmp_path / "venv" + subprocess.run( + [sys.executable, "-m", "venv", "--system-site-packages", str(venv_dir)], + cwd=project_dir, + check=True, + capture_output=True, + text=True, + ) + + python_in_venv = venv_dir / "bin" / "python" + subprocess.run( + [str(python_in_venv), "-m", "pip", "install", "--no-deps", "--force-reinstall", str(wheel_path)], + cwd=project_dir, + check=True, + capture_output=True, + text=True, + ) + + smoke_test = subprocess.run( + [ + str(python_in_venv), + "-c", + ( + "from orca_sim import OrcaHandRight; " + "env = OrcaHandRight(render_mode='rgb_array'); " + "obs, info = env.reset(); " + "print(obs.shape, type(info).__name__); " + "env.close()" + ), + ], + cwd=tmp_path, + check=True, + capture_output=True, + text=True, + ) + + assert "(34,)" in smoke_test.stdout + assert "dict" in smoke_test.stdout