mikasa-robo-vla-medium-envs

Code for 5 "Medium"-horizon MIKASA-Robo-VLA envs that are not (yet) in upstream CognitiveAISystems/MIKASA-Robo, plus a script to reproduce the YCB asset download it needs and two runnable examples (oracle rollout + video, oracle data collection).

Datasets collected from these envs (RLDS and LeRobotDataset v3):

• RLDS: UncertaintyVLA/mikasa-robo-vla-medium-rlds
• LeRobot v3: UncertaintyVLA/mikasa-robo-vla-medium-lerobot

The 5 envs

env_id	max_episode_steps	memory type	language instruction
RememberColor3-Medium-VLA-v0	150	color	Observe the cube's color, wait, then touch the cube of the same color.
RememberColor5-Medium-VLA-v0	150	color	Observe the cube's color, wait, then touch the cube of the same color.
RememberShapeAndColor3x2-Medium-VLA-v0	150	shape+color	Observe the object's shape and color, wait, then touch the object of the same shape and color.
RememberShapeAndColor3x3-Medium-VLA-v0	150	shape+color	Observe the object's shape and color, wait, then touch the object of the same shape and color.
ShellGameColorLampTouch-Medium-VLA-v0	100	spatial	Observe which color is under each cup, then touch the cup matching the lamp color.

They sit at Medium cue/pause durations (CUE_PHASE_STEPS/EMPTY_PHASE_STEPS = [30, 50]) between upstream's Standard ([1, 5]) and Long ([10-100, 50-450]) variants of the same tasks.

Install

pip install mikasa-robo-suite mani-skill   # or your existing MIKASA-Robo install
git clone <this repo>
cd mikasa-robo-vla-medium-envs
pip install -e .
python download_ycb_assets.py              # only needed for ShellGameColorLampTouch-Medium

import gymnasium as gym
import mikasa_robo_vla_medium  # noqa: F401  (registers the 5 env ids on import)

env = gym.make("RememberColor3-Medium-VLA-v0", num_envs=1, obs_mode="state", sim_backend="gpu")
obs, info = env.reset(seed=0)

sim_backend="gpu" is mandatory, not a default to override. All 5 envs' datasets were collected by rolling out a PPO oracle policy (the same 5 checkpoints shipped in oracle_checkpoints/) with sim_backend="gpu". ManiSkill's CPU and GPU PhysX backends are not interchangeable — they diverge in contact behavior, and a checkpoint trained/collected on one backend does not transfer to the other. Evaluating these envs with sim_backend="cpu" will silently give physically different rollouts, not an error.

Evaluation: applying the standard VLA wrapper stack

These envs plug into mikasa_robo_suite's own wrapper-stack machinery — importing this package registers their VLAWrapperConfig entries (curriculum no-op wrapper during cue/empty phases, the render overlays that draw the "Target: ..." text) into the same VLA_WRAPPER_CONFIGS table the stock *-VLA-v0/*-Long-VLA-v0 envs use. So the normal way to build one of these envs for a VLA policy (RGB observations, not the state-only oracle setup in example_rollout.py) is apply_mikasa_vla_wrappers, unmodified:

import gymnasium as gym
import mikasa_robo_vla_medium  # noqa: F401  (registers env ids + wrapper configs)
from mikasa_robo_suite.vla.utils.apply_wrappers import apply_mikasa_vla_wrappers

env = gym.make(
    "RememberColor5-Medium-VLA-v0",
    num_envs=1,
    obs_mode="rgb",
    control_mode="pd_ee_delta_pose",
    render_mode="rgb_array",
    sim_backend="gpu",           # mandatory, see above
    reward_mode="normalized_dense",
)
env = apply_mikasa_vla_wrappers(env)

obs, info = env.reset(seed=0)
# obs["rgb"]:    (num_envs, 128, 128, 6) -- base camera in channels 0:3, wrist in 3:6
# obs["proprio"]: (num_envs, 7) -- converted end-effector pose (xyz + rpy + gripper)
# info["language_instruction"]: per-episode instruction string

This is the same wrapper-application pattern used to evaluate MIKASA-Robo-VLA envs in new_mem_vla (apply_mikasa_vla_wrappers there is manually unrolled per env instead of called directly, because that project reads the wrapper table from an installed mikasa_robo_suite rather than depending on it as a package — the resulting chain is identical: StateOnlyTensorToDictWrapper → curriculum wrapper (if any) → per-task overlays → FlattenRGBDObservationWrapper(rgb=True, depth=False, state=False, oracle=False, joints=True) → ConvertJointsToEEFXyzRpyGripperWrapper). No info field reaches the policy through this chain — the overlay wrappers only affect env.render() for video, not obs.

Why a separate package instead of a PR to MIKASA-Robo

Two of these envs' base classes have known bugs in the upstream repo:

• RememberColorVLABaseEnv/RememberShapeAndColorVLABaseEnv._initialize_episode never set self.prompt, so env.prompt raised AttributeError after reset.
• ShellGameColorLampTouchVLABaseEnv placed the three colored balls in a fixed slot order (red=left, green=center, blue=right) every episode, so the task was solvable from mug position alone, without ever looking at the cue phase.

Rather than patch those classes in place (which would silently change the behavior of the existing *-VLA-v0/*-Long-VLA-v0 envs for anyone already depending on them), every fix here lives in a subclass in envs.py: _initialize_episode/evaluate/_ensure_phase_buffers are overridden, the installed mikasa_robo_suite package is never modified. See envs.py's module docstring for the exact diff. This was a deliberate choice, not an oversight — no upstream issue or PR was filed for it.

One consequence worth knowing if you compare against upstream's own published shell_game_color_lamp_touch_vla_v0 dataset: it was collected from the un-fixed env, so its ball order is not randomized. Ours is.

download_ycb_assets.py: two asset gaps ShellGameColorLampTouch-Medium needs fixed

ShellGameColorLampTouch-Medium-VLA-v0 needs two things this script fixes, both idempotent (they do nothing if already fixed):

1. The YCB mug asset set. ManiSkill's own automatic downloader is broken as of this writing (its bundled checksum no longer matches upstream main), so this pins a known-good commit instead. Non-interactive (the stock path prompts y/n and crashes with EOFError when run in a job/script).
2. low_poly_light_bulb.glb, the lamp shape. The PyPI 1.0.0 wheel of mikasa-robo-suite — the package pip install mikasa-robo-suite above installs — does not include this file, even though the installed source references it. Every ShellGame*Lamp* env, including this one, fails at env creation with RuntimeError: filesystem error: cannot make canonical path pointing at that exact missing path. Confirmed independently by installing into a brand new venv with nothing but the commands in "Install" above. This script copies this package's own bundled copy of the file into the installed mikasa_robo_suite. If you installed mikasa_robo_suite from the pinned git tag instead of PyPI (pip install git+https://github.com/CognitiveAISystems/MIKASA-Robo.git@v1.0.0), the file is already there and this step is a no-op.

Run once, after both pip install steps above:

python download_ycb_assets.py

Examples

python example_rollout.py --env-id ShellGameColorLampTouch-Medium-VLA-v0 --out demo.mp4
python example_collect.py --env-id RememberColor5-Medium-VLA-v0 --episodes 20 --out episodes.npz

Both load one of the 5 oracle_checkpoints/<env_id>.pt PPO checkpoints (state-based, obs_mode="state" only — these are not vision policies). example_collect.py is a minimal single-file collector; it does not reproduce the retry/escape-seed handling the published datasets were actually collected with for seeds the oracle can't solve (see the dataset repos' READMEs).