Wrappers API#

For usage patterns and composition guidance, see the Wrappers Cookbook. This page is a full API reference generated from the source docstrings.

The 19 wrappers are organised into four groups, plus a one-call helper that wires them up for any of the 90 MIKASA-Robo-VLA tasks:

  • One-call helper — picks the correct chain for every task.

  • Core — always required: dict conversion and proprioception remapping.

  • Action-shaping — control robot behaviour during specific episode phases.

  • Task-specific info — inject ground-truth fields into the info dict.

  • Render / debug — overlay diagnostic information on rendered frames.

One-Call Helper#

For every of the 90 MIKASA-Robo-VLA env IDs, use apply_mikasa_vla_wrappers() immediately after gym.make. It applies the canonical wrapper chain that matches the dataset-collection pipeline. The individual wrapper classes below remain available for users that need to compose chains by hand.

apply_mikasa_vla_wrappers(env: Env, *, include_overlays: bool = True) → Env[source]#

Apply the configured MIKASA-Robo-VLA wrapper stack.

Call this immediately after gym.make. Add support for another env by adding its VLAWrapperConfig to VLA_WRAPPER_CONFIGS.

The returned env exposes env.max_episode_steps with the configured ManiSkill horizon, including a gym.make(..., max_episode_steps=...) override.

MIKASA_VLA_ENV_IDS#

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.

Core Wrappers#

class StateOnlyTensorToDictWrapper(env)[source]#

Bases: ObservationWrapper

Wrapper that converts tensor observation to a dictionary with ‘state’ key.

observation(obs)[source]#

Returns a modified observation.

Parameters:

observation – The env observation

Returns:

The modified observation

class ConvertJointsToEEFXyzRpyGripperWrapper(env, qpos_dim: int | None = None, gripper_finger_dims: int = 2)[source]#

Bases: ObservationWrapper

Convert flattened joint-state input into observation[‘proprio’].

The VLA-facing proprio vector is xyz(3) + rpy(3) + gripper(1).

Expected source layout for flattened joints is: [tcp_pose(7), qpos(n), qvel(n), …], where tcp_pose is [x, y, z, qw, qx, qy, qz].

observation(obs)[source]#

Returns a modified observation.

Parameters:

observation – The env observation

Returns:

The modified observation

Action-Shaping Wrappers#

class InitialZeroActionWrapper(env, n_initial_steps=1)[source]#

Bases: ActionWrapper

action(action)[source]#

Modifies action before sending it to the environment

reset(**kwargs)[source]#

Resets the step counter

step(action)[source]#

Runs the env env.step() using the modified action from self.action().

class CurriculumPhaseNoopActionWrapper(env)[source]#

Bases: ActionWrapper

action(action)[source]#

Returns a modified action before env.step() is called.

Parameters:

action – The original step() actions

Returns:

The modified actions

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Runs the env env.step() using the modified action from self.action().

class CurriculumPhaseNoopActionWrapperPdJointPos(env)[source]#

Bases: CurriculumPhaseNoopActionWrapper

Curriculum-phase noop wrapper for envs running in pd_joint_pos control mode.

Plain CurriculumPhaseNoopActionWrapper sends action = 0, which in pd_joint_pos would command the robot to move toward qpos = [0, …, 0] instead of holding the current pose. This subclass overrides the noop action to be the robot’s current arm qpos plus a normalized gripper command — i.e., “stay where you are”.

GRIPPER_HIGH = 0.04#
GRIPPER_LOW = -0.01#
action(action)[source]#

Returns a modified action before env.step() is called.

Parameters:

action – The original step() actions

Returns:

The modified actions

class CameraShutdownWrapper(env, n_initial_steps=19)[source]#

Bases: Wrapper

Wrapper that zeros out all camera observations

if n_initial_steps = 4 then t in [0, 4] (5 steps) action is zero if n_initial_steps = 9 then t in [0, 9] (10 steps) action is zero if n_initial_steps = 19 then t in [0, 19] (20 steps) action is zero

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

Task-Specific Info Wrappers#

class RememberColorInfoWrapper(env)[source]#

Bases: Wrapper

Render the target color as a color swatch for color-memory tasks.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RememberShapeInfoWrapper(env)[source]#

Bases: Wrapper

Render target shape for remember-shape tasks.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RememberShapeAndColorInfoWrapper(env)[source]#

Bases: Wrapper

Render target shape+color text for remember-shape-and-color tasks.

decode_shape(shape_id: int) → str[source]#
render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class MemoryCapacityInfoWrapper(env)[source]#

Bases: Wrapper

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

Render / Debug Wrappers#

class RenderStepInfoWrapper(env)[source]#

Bases: Wrapper

A wrapper that renders the current step count and target cup on the screen.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RenderRewardInfoWrapper(env)[source]#

Bases: Wrapper

A wrapper that renders the current reward on the screen.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RenderPressProgressInfoWrapper(env)[source]#

Bases: Wrapper

Renders button press progress:

  • raw presses (instant physical presses)

  • confirmed presses (completed press cycles)

  • target presses

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RenderWorkingBatteriesInfoWrapper(env)[source]#

Bases: Wrapper

Renders progress of discovered working batteries: found / target.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class ShellGameRenderCupInfoWrapper(env)[source]#

Bases: Wrapper

A wrapper that renders the current step count and target cup on the screen.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RotateRenderAngleInfoWrapper(env)[source]#

Bases: Wrapper

A wrapper that renders the current step count and target cup on the screen.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RenderTraceShapeDebugWrapper(env, minimap_size=160, minimap_top=130)[source]#

Bases: Wrapper

Debug overlay for TraceShape and TraceShapeSeq tasks.

SHAPE_NAMES = {0: 'Circle', 1: 'Square', 2: 'Triangle'}#
render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class RenderTimedTransferInfoWrapper(env)[source]#

Bases: Wrapper

Renders a countdown and timing info for TimedTransfer tasks.

Displayed on each frame:

  • Steps remaining until the target placement moment

  • Current phase (WAIT / COUNTING / WINDOW / LATE)

  • Window bounds and cube-on-red status

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

class DebugRewardWrapper(env)[source]#

Bases: Wrapper

A wrapper that renders the current step count and target cup on the screen.

render()[source]#

Uses the render() of the env that can be overwritten to change the returned data.

reset(**kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#

Uses the step() of the env that can be overwritten to change the returned data.

On this page
Show Source