gym-aloha

The gym-aloha package provides two main tasks:

1. TransferCubeTask: The right arm picks up a red cube from the table and places it into the gripper of the left arm.
2. InsertionTask: The left and right arms pick up a socket and a peg respectively, then perform an insertion task in mid-air.

Success is defined as achieving a maximum reward of 4 points.

The following example demonstrates how to initialize an environment, run a random policy loop, render frames, and save the result as an MP4 video using imageio.

# example.py
import imageio
import gymnasium as gym
import numpy as np
import gym_aloha

env = gym.make("gym_aloha/AlohaInsertion-v0")
observation, info = env.reset()
frames = []

for _ in range(1000):
    action = env.action_space.sample()
    observation, reward, terminated, truncated, info = env.step(action)
    image = env.render()
    frames.append(image)

    if terminated or truncated:
        observation, info = env.reset()

env.close()
imageio.mimsave("example.mp4", np.stack(frames), fps=25)

To install gym-aloha, it is recommended to use a Python 3.10 virtual environment. You can use conda to create and activate the environment before installing via pip.

conda create -y -n aloha python=3.10 && conda activate aloha
pip install gym-aloha

The structure of the returned observation depends on the obs_type provided during initialization:

1. pixels:

   • A gymnasium.spaces.Dict containing:
     • top: A Box array of shape (observation_height, observation_width, 3) with uint8 values (0-255).

2. pixels_agent_pos:

   • A gymnasium.spaces.Dict containing:
     • pixels: A nested Dict with the top camera image (as described above).
     • agent_pos: A Box array of shape (len(JOINTS),) containing joint positions as float64.

3. state:

   • Not currently implemented; calling this will raise a NotImplementedError.

When creating an environment with gym.make(), you can pass the following arguments:

import gymnasium as gym
import gym_aloha

env = gym.make("gym_aloha/AlohaInsertion-v0", obs_type="pixels", render_mode="rgb_array")

The action space is a 14-dimensional continuous vector representing the state of both arms and grippers:

• 6 values per arm: Absolute joint positions.
• 1 value per gripper: Position normalized between 0 (closed) and 1 (open).

Observations are returned as a dictionary containing:

• qpos: Position data for arms and grippers.
• qvel: Velocity data for arms and grippers.
• images: Camera feeds from various angles.
• env_state: Additional environment state (e.g., positions of the peg and sockets).

Use sample_insertion_pose to generate random 7D poses for both a peg and a socket, typically used for insertion tasks. The function returns a tuple of two numpy arrays: (peg_pose, socket_pose).

Each pose is a concatenation of a 3D position [x, y, z] and a 4D quaternion [w, x, y, z] (fixed at [1, 0, 0, 0]).

Peg Position Ranges:

• X: [0.1, 0.2]
• Y: [0.4, 0.6]
• Z: [0.05, 0.05]

Socket Position Ranges:

• X: [-0.2, -0.1]
• Y: [0.4, 0.6]
• Z: [0.05, 0.05]

Args: seed (int, optional): Random seed for reproducibility.

from gym_aloha.utils import sample_insertion_pose

# Sample poses for both peg and socket
peg_pose, socket_pose = sample_insertion_pose(seed=42)

Use sample_box_pose to generate a random 7D pose for a box. The returned pose is a concatenation of a 3D position [x, y, z] and a 4D quaternion [w, x, y, z].

By default, the position is sampled from the following ranges:

• X: [0.0, 0.2]
• Y: [0.4, 0.6]
• Z: [0.05, 0.05]

The quaternion is fixed at [1, 0, 0, 0] (no rotation).

Args: seed (int, optional): Random seed for reproducibility.

from gym_aloha.utils import sample_box_pose

# Sample a random box pose
pose = sample_box_pose(seed=42)
# pose is a numpy array: [x, y, z, qw, qx, qy, qz]

The reset method restarts the environment to a starting state. It accepts an optional seed for reproducibility.

Returns:

• observation: The initial observation formatted by obs_type.
• info (dict): Metadata dictionary, initialized with {"is_success": False}.

Seeding: If a seed is provided, it is used to seed the underlying task's random number generator and the NumPy random state.

observation, info = env.reset(seed=42)