Description

A suite of 3D robotic environments built on OGBench and MuJoCo. The environments feature a UR5e robotic arm with a Robotiq gripper performing various manipulation tasks including pick-and-place, stacking, and complex scene interactions. It also includes a suite of navigation environments.

Three environment families are available:

• Cube: Pure cube manipulation tasks with 1-8 colored cubes
• Scene: Complex scene with cubes, buttons, a drawer, and a window
• PointMaze: 2D maze navigation with point-mass, ant, or humanoid agents

import stable_worldmodel as swm

# Cube environment - set env_type to: single, double, triple, quadruple, or octuple
world = swm.World('swm/OGBCube-v0', num_envs=4, env_type='single')

# Scene environment (cube + drawer + window + buttons)
world = swm.World('swm/OGBScene-v0', num_envs=4)

# PointMaze environment - point agent in a large maze, pixel observations
world = swm.World('swm/OGBPointMaze-v0', num_envs=4)

Cube Environment

Manipulation tasks involving 1-8 colored cubes that must be moved to target positions. Tasks include pick-and-place, swapping positions, and multi-cube stacking.

Success criteria: All cubes must be within 4cm of their target positions.

Environment Specs

Environment Types

Info Dictionary

Variation Space

Scene Environment

A complex manipulation scene with a single cube, two lock buttons, a drawer, and a window. Tasks require interacting with multiple objects in sequence (e.g., press button to unlock drawer, then open drawer).

Success criteria: All objects must be in their target configurations (cube position, button states, drawer position, window position).

Environment Specs

Scene Objects

Tasks

1. Open: Open both drawer and window
2. Unlock and Lock: Toggle locks and close drawer/window
3. Rearrange Medium: Move cube and toggle some objects
4. Put in Drawer: Open drawer and place cube inside
5. Rearrange Hard: Complex multi-step rearrangement

Variation Space

Multiview Mode

The Cube and Scene environments support rendering from multiple camera angles simultaneously:

# Enable multiview rendering
world = swm.World('swm/OGBCube-v0', num_envs=4, env_type='double', multiview=True)

# Returns dict with 'front_pixels' and 'side_pixels' keys
obs = world.reset()

Expert Policy

The Cube and Scene environments include a built-in expert policy for data collection:

from stable_worldmodel.envs.ogbench import ExpertPolicy

policy = ExpertPolicy()
world.set_policy(policy)

PointMaze Environment

2D maze navigation tasks where an agent must reach a goal position (or push a ball to a goal). Built on OGBench's locomaze suite, wrapping it with a variation_space for visual domain randomization.

Success criteria: Agent (or ball) is within a threshold distance of the goal position.

Environment Specs

Configuration

# Default: point agent, large maze, pixel observations
world = swm.World('swm/OGBPointMaze-v0', num_envs=4)

# Custom configuration
world = swm.World(
    'swm/OGBPointMaze-v0',
    num_envs=4,
    loco_env_type='point',   # 'point', 'ant', or 'humanoid'
    maze_env_type='maze',    # 'maze' (navigate) or 'ball' (push ball to goal)
    maze_type='large',       # 'arena', 'medium', 'large', 'giant', 'teleport'
    ob_type='pixels',        # 'pixels' or 'states'
)

Agent Types

Maze Layouts

Maze Task Types

Variation Space

Note: In ob_type='pixels' mode, floor.color only takes effect when floor.pixel_encoding=0. When enabled (=1), OGBench overwrites the floor with a position-encoding gradient (R ← row, G ← column, B = 128) to help the agent infer its location.

Datasets