env.py

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""EnvHub environment: the ROBOTIS OMX multicubes setup of the ``maximellerbach/omx_multicubes`` dataset.

A ROBOTIS OMX arm (ROBOTIS's ``robotis_omx`` MJCF, bundled, Apache-2.0) on a grey board fenced by wooden
slats, with four painted wooden cubes to place in a blue tape square, seen by a top camera and a wrist
camera. The layout, both cameras, the joint zero offsets and the camera delay were fitted to the dataset's
images and joint states.

State and actions use the dataset's units, those of LeRobot's ``omx_follower`` with its default
``use_degrees=False`` and factory calibration: -100..100 for the arm joints (100 = 180 degrees from the
servo center) and 0..100 for the gripper (50 = closed). Actions are absolute position targets at 30 fps.
Each step's reward is the fraction of cubes resting in the square; the episode succeeds (and terminates)
when all four are in, and is truncated after ``episode_length`` steps.
"""

from collections import deque
from functools import partial
from pathlib import Path
from typing import Any

import gymnasium as gym
import mujoco
import numpy as np
from gymnasium import spaces

REPO_ID = "maximellerbach/omx-multicubes-mujoco"
# Not resolved: in the Hugging Face cache, repository files are symlinks into blobs/ named by hash.
HERE = Path(__file__).absolute().parent
# Bundled from ROBOTIS-GIT/robotis_mujoco_menagerie at commit d8344c0dbe7a00208d0301111523dde65efc174a.
ROBOT_MODEL = "robotis_omx/omx.xml"
SCENE_FILE = "omx_multicubes.xml"
MOUNT_SITE = "omx_mount"

JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper")
SIM_JOINTS = ("Joint1", "Joint2", "Joint3", "Joint4", "Joint5", "Gripper")
JOINT_NAMES = [f"{joint}.pos" for joint in JOINTS]
# Every dataset episode starts here: shoulder folded back, forearm forward, gripper open.
HOME = (0.0, -50.0, 50.0, 0.0, 0.0, 60.0)
# Fitted on the dataset: the arm's joint angles minus the model's at the same state (radians). The top
# camera fit gives all five; the shoulder gets 4 more degrees from the grasps (the real arm reaches lower
# than its encoders report, through gear backlash and sag), so that replayed dataset grasps close at
# mid-cube height.
JOINT_OFFSETS = (-0.0157, 0.0920, 0.1046, 0.1542, 0.0135)
# The wrist camera on link5 (in the frame where x points along the gripper), fitted on the dataset.
WRIST_CAMERA = {
    "pos": [0.0440, -0.0039, 0.0350],
    "quat": [0.58949, -0.35461, -0.39399, 0.60953],
    "fovy": 77.17,
}
# The real fingers are Fin Ray style: they bend around what they grasp. Holding a cube, the gripper reads
# 52.8 (10 degrees open) in the datasets, closer than the rigid fingers' geometry allows. Each finger's
# collision hull is cut by the plane its inner face reaches at that opening, so the cut faces close
# parallel on the cube there, as the bent real fingers do, instead of pressing into it.
GRIPPER_HOLD = 52.8
# The gripper servo runs in current-limited position mode: an XL330's stall torque (N.m). The armature
# and friction let it follow the datasets' gripper commands as closely as the real one does.
GRIPPER_FORCE = 0.5
GRIPPER_ARMATURE = 0.005
GRIPPER_FRICTIONLOSS = 0.02
FINGER_FRICTION = (1.5, 0.02, 0.002)
# Stiff finger contacts, as for the cubes in the scene: soft ones let the cubes sink and slip.
FINGER_SOLREF = (0.004, 1.0)
FINGER_SOLIMP = (0.95, 0.99, 0.001, 0.5, 2.0)
# The dataset's images lag its joint states by about two frames.
IMAGE_DELAY = 2

# How the LeRobot simulator server describes this environment (ignored in-process).
LEROBOT_SIM = {
    "fps": 30,
    "control": "position",
    "semantics": "omx-follower-range-m100-v1",
    "state_names": JOINT_NAMES,
    "action_names": JOINT_NAMES,
    "gripper_indices": [5],
}

CUBE_BODIES = ("red_cube", "green_cube", "blue_cube", "yellow_cube")
# 2.9 cm cubes, fitted on the dataset's top images.
CUBE_HALF = 0.0145
# One suite per dataset: its cubes, the x, y ranges their centers start in, instruction, episode length.
TASKS: dict[str, dict[str, Any]] = {
    # maximellerbach/omx_multicubes: four cubes in a loose cluster on the arm's left.
    "omx_multicubes": {
        "cubes": CUBE_BODIES,
        "region": ((0.09, 0.22), (0.14, 0.28)),
        "instruction": "pick all the cubes and place them one by one in the blue square",
        "episode_length": 3600,
    },
    # maximellerbach/omx_pickandplace: the red cube alone, anywhere on the arm's left.
    "omx_pickandplace": {
        "cubes": ("red_cube",),
        "region": ((0.15, 0.33), (0.04, 0.26)),
        "instruction": "Pick the red cube and place it in the blue square",
        "episode_length": 900,
    },
}
# The tape square (outer edges); a cube is in it when its center is inside and it rests on the board.
SQUARE = ((0.108, 0.236), (-0.264, -0.111))
IN_SQUARE_MAX_HEIGHT = 0.03


def repo_root() -> Path:
    """Directory holding the scene and the arm model.

    LeRobot's in-process loader downloads ``env.py`` alone, into ``snapshots/<commit>/`` of the Hub
    cache; fetch the rest of that same commit. A local checkout or a full snapshot needs nothing.
    """
    if (HERE / SCENE_FILE).is_file():
        return HERE
    from huggingface_hub import snapshot_download

    revision = HERE.name if HERE.parent.name == "snapshots" else None
    return Path(snapshot_download(REPO_ID, revision=revision))


def finger_contact_vertices(arm: mujoco.MjModel) -> dict[str, np.ndarray]:
    """Each finger's collision mesh vertices in its body frame, cut by the plane of ``GRIPPER_HOLD``.

    At that opening the two cut faces are parallel to the gripper's axis, a cube's width apart; the
    vertices beyond a finger's plane are moved onto it, so the finger's convex hull ends at the plane.
    """
    hold = (GRIPPER_HOLD - 50) * np.pi / 50
    pivots = {finger: arm.body(finger).pos for finger in ("link6", "link7")}
    center = (pivots["link6"][1] + pivots["link7"][1]) / 2
    vertices = {}
    # link6 opens toward +y of link5, its mimic link7 toward -y.
    for finger, side in (("link6", 1), ("link7", -1)):
        body = arm.body(finger).id
        geom = next(g for g in range(arm.ngeom) if arm.geom_bodyid[g] == body and arm.geom_contype[g])
        mesh = arm.geom_dataid[geom]
        start = arm.mesh_vertadr[mesh]
        rotation = np.zeros(9)
        mujoco.mju_quat2Mat(rotation, arm.geom_quat[geom])
        local = (
            arm.geom_pos[geom]
            + arm.mesh_vert[start : start + arm.mesh_vertnum[mesh]] @ rotation.reshape(3, 3).T
        )
        # The finger at the holding opening, in link5's frame.
        c, s = np.cos(side * hold), np.sin(side * hold)
        opened = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]])
        held = pivots[finger] + local @ opened.T
        plane = center + side * CUBE_HALF
        held[:, 1] = np.maximum(held[:, 1], plane) if side > 0 else np.minimum(held[:, 1], plane)
        vertices[finger] = (held - pivots[finger]) @ opened
    return vertices


def build_model(root: Path) -> mujoco.MjModel:
    """Compile the scene with the OMX attached at its mount site, its gripper and the wrist camera on link5."""
    spec = mujoco.MjSpec.from_file(str(root / SCENE_FILE))
    arm = mujoco.MjSpec.from_file(str(root / ROBOT_MODEL))
    # The arm's keyframe would not match the scene's free cubes.
    for key in list(arm.keys):
        arm.delete(key)
    arm.body("link5").add_camera(name="wrist", **WRIST_CAMERA)
    for joint in ("Gripper", "Gripper_mimic"):
        arm.joint(joint).armature = GRIPPER_ARMATURE
        arm.joint(joint).frictionloss = GRIPPER_FRICTIONLOSS
    arm.actuator("Gripper").forcerange = [-GRIPPER_FORCE, GRIPPER_FORCE]
    contact = finger_contact_vertices(mujoco.MjSpec.from_file(str(root / ROBOT_MODEL)).compile())
    for finger, vertices in contact.items():
        arm.add_mesh(name=f"{finger}_contact", uservert=vertices.flatten().tolist())
        for geom in arm.body(finger).geoms:
            if geom.contype or geom.conaffinity:
                geom.meshname = f"{finger}_contact"
                geom.pos, geom.quat = [0, 0, 0], [1, 0, 0, 0]
                geom.friction = FINGER_FRICTION
                geom.condim = 4
                geom.solref = FINGER_SOLREF
                geom.solimp = FINGER_SOLIMP
    spec.attach(arm, site=MOUNT_SITE, prefix="")
    return spec.compile()


class _World:
    """Compiled model, physics state, renderer and reset randomization of the OMX scene."""

    def __init__(self, root: Path, fps: int, height: int, width: int):
        self.model = build_model(root)
        # The control period must be a whole number of physics steps.
        self.substeps = max(1, round(1 / (fps * self.model.opt.timestep)))
        self.model.opt.timestep = 1 / (fps * self.substeps)
        self.data = mujoco.MjData(self.model)
        self.renderer = mujoco.Renderer(self.model, height, width)
        self.qadr = np.array([self.model.joint(name).qposadr[0] for name in SIM_JOINTS])
        self.mimic = self.model.joint("Gripper_mimic").qposadr[0]
        self.actuators = np.array([self.model.actuator(name).id for name in SIM_JOINTS])
        self.lo, self.hi = self.model.actuator_ctrlrange[self.actuators].T.copy()
        self.cubes = {name: self.model.body(name).id for name in CUBE_BODIES}
        self.cube_qadr = {
            name: self.model.jnt_qposadr[self.model.body_jntadr[body]] for name, body in self.cubes.items()
        }

    def reset(
        self, rng: np.random.Generator, joints: np.ndarray, joint_noise: float, cubes: tuple[str, ...], region
    ) -> None:
        """Put the arm at its initial joints, the task's cubes at seeded, non-overlapping poses in ``region``
        and the other cubes off the board."""
        model, data = self.model, self.data
        mujoco.mj_resetData(model, data)
        q = self.to_ctrl(joints + rng.uniform(-joint_noise, joint_noise, len(JOINTS)))
        data.qpos[self.qadr] = q
        data.qpos[self.mimic] = -q[5]
        data.ctrl[self.actuators] = q
        for k, name in enumerate(n for n in CUBE_BODIES if n not in cubes):
            adr = self.cube_qadr[name]
            data.qpos[adr : adr + 3] = (1.0 + 0.1 * k, 1.0, CUBE_HALF - 0.02)
        placed: list[np.ndarray] = []
        for name in cubes:
            adr = self.cube_qadr[name]
            for _ in range(200):
                xy = np.array([rng.uniform(*region[0]), rng.uniform(*region[1])])
                if all(np.linalg.norm(xy - other) > 2.9 * CUBE_HALF for other in placed):
                    break
            else:
                raise ValueError("Could not place the cubes without overlap")
            placed.append(xy)
            yaw = rng.uniform(-np.pi, np.pi)
            data.qpos[adr : adr + 3] = (*xy, CUBE_HALF)
            data.qpos[adr + 3 : adr + 7] = (np.cos(yaw / 2), 0.0, 0.0, np.sin(yaw / 2))
        mujoco.mj_forward(model, data)

    def to_ctrl(self, values: np.ndarray) -> np.ndarray:
        """Convert the dataset's units to joint radians."""
        values = values.astype(np.float64)
        ctrl = np.empty(len(JOINTS))
        ctrl[:5] = values[:5] * np.pi / 100 + JOINT_OFFSETS
        ctrl[5] = (values[5] - 50) * np.pi / 50
        return np.clip(ctrl, self.lo, self.hi)

    def joint_state(self) -> np.ndarray:
        """Joint positions in the units of ``to_ctrl``."""
        q = self.data.qpos[self.qadr]
        state = np.empty(len(JOINTS))
        state[:5] = (q[:5] - JOINT_OFFSETS) * 100 / np.pi
        state[5] = q[5] * 50 / np.pi + 50
        return state.astype(np.float32)

    def cubes_in_square(self, cubes: tuple[str, ...]) -> int:
        """Number of ``cubes`` resting with their center in the tape square."""
        xyz = self.data.xpos[[self.cubes[name] for name in cubes]]
        inside = (
            (xyz[:, 0] > SQUARE[0][0])
            & (xyz[:, 0] < SQUARE[0][1])
            & (xyz[:, 1] > SQUARE[1][0])
            & (xyz[:, 1] < SQUARE[1][1])
            & (xyz[:, 2] < IN_SQUARE_MAX_HEIGHT)
        )
        return int(inside.sum())

    def render(self, camera: str) -> np.ndarray:
        self.renderer.update_scene(self.data, camera=camera)
        return self.renderer.render().copy()


class OmxMulticubesEnv(gym.Env):
    """OMX joint position control on the multicubes board, observed through the top and wrist cameras."""

    metadata = {"render_modes": ["rgb_array"], "render_fps": 30}

    def __init__(
        self,
        task: str = "omx_multicubes",
        fps: int = 30,
        episode_length: int | None = None,
        observation_height: int = 480,
        observation_width: int = 640,
        instruction: str | None = None,
        initial_joints: tuple[float, ...] = HOME,
        joint_noise: float = 1.0,
        image_delay: int = IMAGE_DELAY,
    ):
        super().__init__()
        if task not in TASKS:
            raise ValueError(f"Unknown task {task!r}; choose from {sorted(TASKS)}")
        self.task = task
        self.cube_names = TASKS[task]["cubes"]
        self.fps = fps
        self._max_episode_steps = episode_length or TASKS[task]["episode_length"]
        self._steps = 0
        self.height = observation_height
        self.width = observation_width
        self.instruction = instruction
        self.initial_joints = np.asarray(initial_joints, dtype=np.float32)
        if self.initial_joints.shape != (len(JOINTS),):
            raise ValueError(f"initial_joints needs one value per joint {JOINTS}")
        self.joint_noise = joint_noise
        # Observations pair the current joint state with the images of `image_delay` steps earlier.
        self.image_delay = image_delay
        self._images: deque[dict[str, np.ndarray]] = deque(maxlen=image_delay + 1)
        self.cameras = ("top", "wrist")
        self._world: _World | None = None

        image = spaces.Box(0, 255, (self.height, self.width, 3), np.uint8)
        self.observation_space = spaces.Dict(
            {
                "agent_pos": spaces.Box(-200.0, 200.0, (len(JOINTS),), np.float32),
                "pixels": spaces.Dict(dict.fromkeys(self.cameras, image)),
            }
        )
        self.action_space = spaces.Box(-200.0, 200.0, (len(JOINTS),), np.float32)

    @property
    def task_description(self) -> str:
        """The configured instruction, else the dataset's task."""
        return self.instruction if self.instruction is not None else TASKS[self.task]["instruction"]

    def _ensure_world(self) -> _World:
        """Build the world on first use, so async workers each own their GL context."""
        if self._world is None:
            self._world = _World(repo_root(), self.fps, self.height, self.width)
        return self._world

    def render(self, camera: str | None = None) -> np.ndarray:
        """Render one camera now (no delay), the top camera by default."""
        return self._ensure_world().render(camera or "top")

    def _obs(self, world: _World) -> dict[str, Any]:
        self._images.append({camera: world.render(camera) for camera in self.cameras})
        return {"agent_pos": world.joint_state(), "pixels": self._images[0]}

    def reset(self, seed: int | None = None, **kwargs) -> tuple[dict[str, Any], dict[str, Any]]:
        """Put the arm at its initial pose and the cubes at seeded positions."""
        world = self._ensure_world()
        super().reset(seed=seed)
        world.reset(
            self.np_random, self.initial_joints, self.joint_noise, self.cube_names, TASKS[self.task]["region"]
        )
        self._steps = 0
        self._images.clear()
        first = {camera: world.render(camera) for camera in self.cameras}
        self._images.extend([first] * self.image_delay)
        return self._obs(world), {"is_success": False, "cubes_in_square": 0}

    def step(self, action: np.ndarray) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
        """Hold the joint targets for one control period."""
        world = self._ensure_world()
        action = np.asarray(action)
        if action.shape != self.action_space.shape:
            raise ValueError(f"Expected action of shape {self.action_space.shape}, got {action.shape}")
        world.data.ctrl[world.actuators] = world.to_ctrl(action)
        mujoco.mj_step(world.model, world.data, nstep=world.substeps)
        self._steps += 1
        in_square = world.cubes_in_square(self.cube_names)
        success = in_square == len(self.cube_names)
        truncated = self._steps >= self._max_episode_steps
        info = {"is_success": success, "cubes_in_square": in_square}
        return self._obs(world), in_square / len(self.cube_names), success, truncated, info

    def close(self) -> None:
        if self._world is not None:
            self._world.renderer.close()
            self._world = None


def make_env(n_envs: int = 1, use_async_envs: bool = False, **kwargs: Any) -> dict[str, dict[int, Any]]:
    """EnvHub entry point: one suite per task, ``{task: {0: vector_env}}``; ``kwargs`` override the env's."""
    vector_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
    return {
        task: {0: vector_cls([partial(OmxMulticubesEnv, task=task, **kwargs)] * n_envs)} for task in TASKS
    }