ocbench 1.0.0__py3-none-any.whl

ocbench/__init__.py

@@ -0,0 +1,440 @@
+"""OCBench manipulation environments."""
+
+import gymnasium
+from gymnasium.envs.registration import register, registry
+from ocbench.dataset_utils import download_datasets, get_episode_offsets, load_dataset
+
+__all__ = (
+    'download_datasets',
+    'get_episode_offsets',
+    'load_dataset',
+    'make',
+    'make_env_and_datasets',
+    'parse_env_spec',
+)
+
+_LITE_KWARGS = dict(
+    action_delta_scale=5.0,
+    lite=True,
+    control_timestep=0.1,
+)
+
+
+def parse_env_spec(env_name):
+    """Return the family, backend, constructor kwargs, and episode limit."""
+    visual = env_name.startswith('visual-')
+    family, task = env_name.removeprefix('visual-').split('-', 1)
+    backend = 'cpu' if task.startswith('cpu-') else 'mjwarp'
+    cpu_name = env_name if backend == 'cpu' else f'{family}-cpu-{task}'
+    if visual and backend == 'mjwarp':
+        cpu_name = f'visual-{cpu_name}'
+    spec = gymnasium.spec(cpu_name)
+    return family, backend, dict(spec.kwargs), int(spec.max_episode_steps)
+
+
+def make(env_name, **kwargs):
+    """Make a batched MJWarp environment, or a single CPU environment with a `-cpu-` ID.
+
+    MJWarp uses 1024 worlds by default; pass `nworld` to choose the batch size.
+    For MJWarp environments, stop each episode after the step limit returned by `parse_env_spec`.
+    CPU environments truncate episodes automatically at their time limit.
+    """
+    family, backend, env_kwargs, _ = parse_env_spec(env_name)
+    if backend == 'cpu':
+        return gymnasium.make(env_name, **kwargs)
+    from ocbench.mjwarp import make_env
+
+    nworld = kwargs.pop('nworld', 1024)
+    env_kwargs.update(kwargs)
+    return make_env(family, env_kwargs, nworld)
+
+
+def make_env_and_datasets(env_name, dataset_root=None, num_shards=None, **env_kwargs):
+    """Make an OCBench environment and load its released datasets.
+
+    Args:
+        env_name: Environment name (e.g., 'block-lite-single-task1-v0').
+        dataset_root: Root directory for downloaded datasets. Defaults to '$XDG_CACHE_HOME/ocbench/datasets' if
+            XDG_CACHE_HOME is set, or '~/.cache/ocbench/datasets' otherwise.
+        num_shards: Number of training shards to load, in filename order. If None, load all shards.
+        **env_kwargs: Keyword arguments to pass to make (e.g., nworld=1024 for MJWarp).
+
+    Returns:
+        A tuple of the environment, training dataset, and validation dataset. Each dataset is a dictionary of
+        NumPy arrays; see load_dataset for the dataset format.
+    """
+    train_paths, val_paths = download_datasets(env_name, dataset_root=dataset_root, num_shards=num_shards)
+    train_dataset = load_dataset(train_paths)
+    val_dataset = load_dataset(val_paths)
+    env = make(env_name, **env_kwargs)
+    return env, train_dataset, val_dataset
+
+
+# Block environments.
+register(
+    id='block-cpu-single-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=1250,
+    kwargs=dict(env_type='single', task_id=1),
+)
+register(
+    id='block-cpu-double-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=2500,
+    kwargs=dict(env_type='double', task_id=1),
+)
+register(
+    id='block-cpu-double-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=2500,
+    kwargs=dict(env_type='double', task_id=2),
+)
+register(
+    id='block-cpu-triple-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=3750,
+    kwargs=dict(env_type='triple', task_id=1),
+)
+register(
+    id='block-cpu-triple-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=3750,
+    kwargs=dict(env_type='triple', task_id=2),
+)
+register(
+    id='block-cpu-quadruple-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=5000,
+    kwargs=dict(env_type='quadruple', task_id=1),
+)
+register(
+    id='block-cpu-quadruple-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=5000,
+    kwargs=dict(env_type='quadruple', task_id=2),
+)
+register(
+    id='block-cpu-quadruple-task3-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=5000,
+    kwargs=dict(env_type='quadruple', task_id=3),
+)
+
+# Chamber environments.
+register(
+    id='chamber-cpu-easy-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=4000,
+    kwargs=dict(env_type='easy', task_id=1),
+)
+register(
+    id='chamber-cpu-easy-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=4000,
+    kwargs=dict(env_type='easy', task_id=2),
+)
+register(
+    id='chamber-cpu-easy-task3-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=4000,
+    kwargs=dict(env_type='easy', task_id=3),
+)
+register(
+    id='chamber-cpu-medium-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=5000,
+    kwargs=dict(env_type='medium', task_id=1),
+)
+register(
+    id='chamber-cpu-medium-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=5000,
+    kwargs=dict(env_type='medium', task_id=2),
+)
+register(
+    id='chamber-cpu-hard-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=5500,
+    kwargs=dict(env_type='hard', task_id=1),
+)
+register(
+    id='chamber-cpu-hard-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=5500,
+    kwargs=dict(env_type='hard', task_id=2),
+)
+
+# Switch environments.
+register(
+    id='switch-cpu-3x3-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=1500,
+    kwargs=dict(env_type='3x3', task_id=1),
+)
+register(
+    id='switch-cpu-3x3-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=1500,
+    kwargs=dict(env_type='3x3', task_id=2),
+)
+register(
+    id='switch-cpu-4x4-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=1500,
+    kwargs=dict(env_type='4x4', task_id=1),
+)
+register(
+    id='switch-cpu-4x4-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=1500,
+    kwargs=dict(env_type='4x4', task_id=2),
+)
+register(
+    id='switch-cpu-5x5-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=4000,
+    kwargs=dict(env_type='5x5', task_id=1),
+)
+register(
+    id='switch-cpu-5x5-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=4000,
+    kwargs=dict(env_type='5x5', task_id=2),
+)
+
+# Hanoi environments.
+register(
+    id='hanoi-cpu-single-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=1000,
+    kwargs=dict(env_type='single', task_id=1),
+)
+register(
+    id='hanoi-cpu-single-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=1000,
+    kwargs=dict(env_type='single', task_id=2),
+)
+register(
+    id='hanoi-cpu-double-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=3000,
+    kwargs=dict(env_type='double', task_id=1),
+)
+register(
+    id='hanoi-cpu-double-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=3000,
+    kwargs=dict(env_type='double', task_id=2),
+)
+register(
+    id='hanoi-cpu-triple-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=7000,
+    kwargs=dict(env_type='triple', task_id=1),
+)
+register(
+    id='hanoi-cpu-triple-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=7000,
+    kwargs=dict(env_type='triple', task_id=2),
+)
+
+# Bowling environment.
+register(
+    id='bowling-cpu-task1-v0',
+    entry_point='ocbench.envs.bowling_env:BowlingEnv',
+    max_episode_steps=500,
+    kwargs=dict(task_id=1),
+)
+
+# Lite block environments.
+register(
+    id='block-cpu-lite-single-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=200,
+    kwargs=dict(env_type='single', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-double-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=400,
+    kwargs=dict(env_type='double', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-double-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=400,
+    kwargs=dict(env_type='double', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-triple-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=600,
+    kwargs=dict(env_type='triple', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-triple-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=600,
+    kwargs=dict(env_type='triple', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-quadruple-task1-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=800,
+    kwargs=dict(env_type='quadruple', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-quadruple-task2-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=800,
+    kwargs=dict(env_type='quadruple', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='block-cpu-lite-quadruple-task3-v0',
+    entry_point='ocbench.envs.block_env:BlockEnv',
+    max_episode_steps=800,
+    kwargs=dict(env_type='quadruple', task_id=3, **_LITE_KWARGS),
+)
+
+# Lite chamber environments.
+register(
+    id='chamber-cpu-lite-easy-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=500,
+    kwargs=dict(env_type='easy', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-easy-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=500,
+    kwargs=dict(env_type='easy', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-easy-task3-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=500,
+    kwargs=dict(env_type='easy', task_id=3, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-medium-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=600,
+    kwargs=dict(env_type='medium', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-medium-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=600,
+    kwargs=dict(env_type='medium', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-hard-task1-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=700,
+    kwargs=dict(env_type='hard', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='chamber-cpu-lite-hard-task2-v0',
+    entry_point='ocbench.envs.chamber_env:ChamberEnv',
+    max_episode_steps=700,
+    kwargs=dict(env_type='hard', task_id=2, **_LITE_KWARGS),
+)
+
+# Lite switch environments.
+register(
+    id='switch-cpu-lite-3x3-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='3x3', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='switch-cpu-lite-3x3-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='3x3', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='switch-cpu-lite-4x4-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='4x4', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='switch-cpu-lite-4x4-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='4x4', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='switch-cpu-lite-5x5-task1-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=800,
+    kwargs=dict(env_type='5x5', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='switch-cpu-lite-5x5-task2-v0',
+    entry_point='ocbench.envs.switch_env:SwitchEnv',
+    max_episode_steps=800,
+    kwargs=dict(env_type='5x5', task_id=2, **_LITE_KWARGS),
+)
+
+# Lite Hanoi environments.
+register(
+    id='hanoi-cpu-lite-single-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='single', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='hanoi-cpu-lite-single-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=300,
+    kwargs=dict(env_type='single', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='hanoi-cpu-lite-double-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=900,
+    kwargs=dict(env_type='double', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='hanoi-cpu-lite-double-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=900,
+    kwargs=dict(env_type='double', task_id=2, **_LITE_KWARGS),
+)
+register(
+    id='hanoi-cpu-lite-triple-task1-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=2100,
+    kwargs=dict(env_type='triple', task_id=1, **_LITE_KWARGS),
+)
+register(
+    id='hanoi-cpu-lite-triple-task2-v0',
+    entry_point='ocbench.envs.hanoi_env:HanoiEnv',
+    max_episode_steps=2100,
+    kwargs=dict(env_type='triple', task_id=2, **_LITE_KWARGS),
+)
+
+# Visual environments.
+for _spec in list(registry.values()):
+    if (
+        isinstance(_spec.entry_point, str)
+        and _spec.entry_point.startswith('ocbench.envs.')
+        and not _spec.id.startswith('visual-')
+    ):
+        register(
+            id=f'visual-{_spec.id}',
+            entry_point=_spec.entry_point,
+            max_episode_steps=_spec.max_episode_steps,
+            kwargs=dict(
+                **_spec.kwargs,
+                ob_type='pixels',
+                width=224,
+                height=224,
+                visualize_info=False,
+                pixel_cameras=('front', 'side', 'ur5e/wrist'),
+            ),
+        )

ocbench/controllers/__init__.py

@@ -0,0 +1,3 @@
+from ocbench.controllers.diff_ik import DiffIKController
+
+__all__ = ('DiffIKController',)

ocbench/controllers/diff_ik.py

@@ -0,0 +1,115 @@
+import mujoco
+import numpy as np
+
+PI = np.pi
+PI_2 = 2 * np.pi
+
+
+def angle_diff(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
+    return np.mod(q1 - q2 + PI, PI_2) - PI
+
+
+class DiffIKController:
+    """Differential inverse kinematics controller."""
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        sites: list,
+        qpos0: np.ndarray = None,
+        damping_coeff: float = 1e-12,
+        max_angle_change: float = np.radians(45),
+    ):
+        self._model = model
+        self._data = mujoco.MjData(self._model)
+        self._qp0 = qpos0
+        self._max_angle_change = max_angle_change
+
+        # Cache references.
+        self._ns = len(sites)  # Number of sites.
+        self._site_ids = np.asarray([self._model.site(s).id for s in sites])
+
+        # Preallocate arrays.
+        self._err = np.empty((self._ns, 6))
+        self._site_quat = np.empty((self._ns, 4))
+        self._site_quat_inv = np.empty((self._ns, 4))
+        self._err_quat = np.empty((self._ns, 4))
+        self._jac = np.empty((6 * self._ns, self._model.nv))
+        self._damping = damping_coeff * np.eye(6 * self._ns)
+        self._eye = np.eye(self._model.nv)
+
+    def _forward_kinematics(self) -> None:
+        """Minimal computation required for forward kinematics."""
+        mujoco.mj_kinematics(self._model, self._data)
+        mujoco.mj_comPos(self._model, self._data)  # Required for mj_jacSite.
+
+    def _integrate(self, update: np.ndarray) -> None:
+        """Integrate the joint velocities in-place."""
+        mujoco.mj_integratePos(self._model, self._data.qpos, update, 1.0)
+
+    def _compute_translational_error(self, pos: np.ndarray) -> None:
+        """Compute the error between the desired and current site positions."""
+        self._err[:, :3] = pos - self._data.site_xpos[self._site_ids]
+
+    def _compute_rotational_error(self, quat: np.ndarray) -> None:
+        """Compute the error between the desired and current site orientations."""
+        for i, site_id in enumerate(self._site_ids):
+            mujoco.mju_mat2Quat(self._site_quat[i], self._data.site_xmat[site_id])
+            mujoco.mju_negQuat(self._site_quat_inv[i], self._site_quat[i])
+            mujoco.mju_mulQuat(self._err_quat[i], quat[i], self._site_quat_inv[i])
+            mujoco.mju_quat2Vel(self._err[i, 3:], self._err_quat[i], 1.0)
+
+    def _compute_jacobian(self) -> None:
+        """Update site end-effector Jacobians."""
+        for i, site_id in enumerate(self._site_ids):
+            jacp = self._jac[6 * i : 6 * i + 3]
+            jacr = self._jac[6 * i + 3 : 6 * i + 6]
+            mujoco.mj_jacSite(self._model, self._data, jacp, jacr, site_id)
+
+    def _error_threshold_reached(self, pos_thresh: float, ori_thresh: float) -> bool:
+        """Return True if position and rotation errors are below the thresholds."""
+        pos_achieved = np.linalg.norm(self._err[:, :3]) <= pos_thresh
+        ori_achieved = np.linalg.norm(self._err[:, 3:]) <= ori_thresh
+        return pos_achieved and ori_achieved
+
+    def _solve(self) -> np.ndarray:
+        """Solve for joint velocities using damped least squares."""
+        H = self._jac @ self._jac.T + self._damping
+        x = self._jac.T @ np.linalg.solve(H, self._err.ravel())
+        if self._qp0 is not None:
+            jac_pinv = np.linalg.pinv(H)
+            q_err = angle_diff(self._qp0, self._data.qpos)
+            x += (self._eye - (self._jac.T @ jac_pinv) @ self._jac) @ q_err
+        return x
+
+    def _scale_update(self, update: np.ndarray) -> np.ndarray:
+        """Scale down update so that the max allowable angle change is not exceeded."""
+        update_max = np.max(np.abs(update))
+        if update_max > self._max_angle_change:
+            update *= self._max_angle_change / update_max
+        return update
+
+    def solve(
+        self,
+        pos: np.ndarray,
+        quat: np.ndarray,
+        curr_qpos: np.ndarray,
+        max_iters: int = 20,
+        pos_thresh: float = 1e-4,
+        ori_thresh: float = 1e-4,
+    ) -> np.ndarray:
+        self._data.qpos = curr_qpos
+
+        for _ in range(max_iters):
+            self._forward_kinematics()
+
+            self._compute_translational_error(np.atleast_2d(pos))
+            self._compute_rotational_error(np.atleast_2d(quat))
+            if self._error_threshold_reached(pos_thresh, ori_thresh):
+                break
+
+            self._compute_jacobian()
+            update = self._scale_update(self._solve())
+            self._integrate(update)
+
+        return self._data.qpos.copy()