tsakman23/visual_masked_distracting_metaworld_sam

Visual Masked Distracting Meta-World (ground-truth + SAM masks)

• —Author: Georgios Tsakoumakis
• —Thesis: Interaction-Masked Latent Action Models for Object-Aware Manipulation under Visual Distractors (MSc, Imperial College London)

Expert Meta-World manipulation trajectories rendered with dynamic video-background distractors, carrying both ground-truth and SAM-predicted segmentation masks for the agent and the manipulated object, plus the manipulated object's pose.

This is the SAM variant of the masked distracting Meta-World data: alongside the exact simulator masks (mask, object_mask) it adds predicted masks (pred_mask, pred_object_mask) obtained by running the Segment Anything Model (SAM) on the frames. Having both lets you train or evaluate under exact vs. imperfect segmentation and measure sensitivity to mask quality. (A ground-truth-only version, with an additional dial-turn-v3 task, lives at `tsakman23/visual_masked_distracting_metaworld`.)

Why this dataset

Distracting Meta-World places a task-irrelevant, moving video background behind the scene, so pixel reconstruction alone no longer isolates the controllable content. The agent mask identifies the robot; the object mask identifies the manipulated object. Together they let a model (or an evaluator) separate embodiment, manipulated object, and distractor background. The dataset was built to support:

• —IM-LAM - a directed agent to object forward-dynamics model over agent + object masks.
• —Foreground-MaskLAM - gating the world-model reconstruction loss by the agent-and-object union so the decoder must reconstruct the manipulated object, not just the arm.

The predicted masks additionally support ablating ground-truth vs. SAM-predicted masks, i.e. how much these models rely on perfect segmentation.

Data fields

Every example is a single timestep t of an expert trajectory:

(observation, action, reward, terminated, truncated) follow the standard (o_t, a_t, r_t, term_t, trunc_t) convention. Frames are stored in trajectory order, so a temporal stack of length k can be formed by slicing consecutive rows.

Configurations (tasks)

The dataset is organized as one HuggingFace builder config per Meta-World task. Load a task by passing its bare config name (no masked-/distracting- prefix), for example push-v3 or handle-pull-v3. Available task configs:

push-v3         sweep-into-v3       door-open-v3
handle-pull-v3  pick-place-v3       peg-insert-side-v3

Splits

Per task configuration:

Usage

from datasets import load_dataset

ds = load_dataset(
    "tsakman23/visual_masked_distracting_metaworld_sam",
    name="handle-pull-v3",   # bare task config name
    split="train",
)

ex = ds[0]
ex["observation_distracted"]  # PIL.Image, 128x128 RGB (distracted frame)
ex["mask"]                    # PIL.Image, 128x128 L, ground-truth agent = 255
ex["object_mask"]             # PIL.Image, 128x128 L, ground-truth object = 255
ex["pred_mask"]               # PIL.Image, 128x128 L, SAM-predicted agent
ex["pred_object_mask"]        # PIL.Image, 128x128 L, SAM-predicted object
ex["object_state"]            # [x, y, z, qw, qx, qy, qz]
ex["action"]                  # [dx, dy, dz, gripper]

To load only the columns you need (masks are large), pass columns=... when constructing a torch-formatted view. Train on ground-truth masks with ["observation_distracted", "mask", "object_mask", ...], or on predicted masks with ["observation_distracted", "pred_mask", "pred_object_mask", ...].

For distracting/masked training you typically use observation_distracted as the observation and drop the vanilla observation column.

How the masks were produced

The ground-truth masks (mask, object_mask) are the simulator's own segmentation, obtained from MuJoCo's segmentation render (per-geom object IDs) at the same camera pose as the RGB frame, then binarized into the agent and manipulated-object masks. object_state is the manipulated object's body pose read from the simulator.

The predicted masks (pred_mask, pred_object_mask) are produced by running the Segment Anything Model (SAM) on the frames and selecting the agent / manipulated-object regions. They approximate the ground-truth masks and, unlike them, can contain segmentation errors - which is exactly what makes them useful for studying robustness to imperfect masks.

Masks

This repository carries both mask sources for every frame: the ground-truth simulator masks (mask, object_mask) and the SAM-predicted masks (pred_mask, pred_object_mask). Use whichever pair suits the experiment - the ground-truth masks for clean supervision/evaluation, the predicted masks to test reliance on segmentation quality.

Licensing and attribution

• —Meta-World environments and scripted expert policies: MIT License (Yu et al., "Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning", CoRL 2019).
• —Distractor backgrounds are video frames from the DAVIS dataset, which is released for non-commercial research use; downstream use of the distracted frames should respect DAVIS's terms.
• —Predicted masks were generated with the Segment Anything Model (SAM) (Kirillov et al., "Segment Anything", ICCV 2023).

This dataset is released under CC BY 4.0 (attribution required). One carve-out: because the observation_distracted frames embed DAVIS video content, that portion is subject to DAVIS's research-only terms - review the DAVIS license before any commercial use or redistribution of the distracted frames. The vanilla observation, masks, object_state, state, and action fields derive from Meta-World (MIT) and its scripted experts.

Citation

This dataset was produced as part of the following MSc thesis:

Georgios Tsakoumakis. Interaction-Masked Latent Action Models for Object-Aware Manipulation under Visual Distractors. MSc thesis, Imperial College London, 2026.

@mastersthesis{tsakoumakis2026imlam,
  title  = {Interaction-Masked Latent Action Models for Object-Aware Manipulation under Visual Distractors},
  author = {Tsakoumakis, Georgios},
  school = {Imperial College London},
  year   = {2026},
  type   = {{MSc} thesis}
}

If you use this dataset, please also cite Meta-World, the DAVIS background dataset, and the Segment Anything Model (SAM).