EpicPinkPenguin/atari_frogger_with_masks

Atari Frogger Expert Trajectories with Player Masks

This dataset contains expert-policy transitions from Atari Frogger together with a programmatically derived binary mask of the player in every stored RGB observation. The expert is a Rainbow agent trained with CleanRL, and player localization is provided by OCAtari.

The dataset accompanies the paper *Segment to Focus: Guiding Latent Action Models in the Presence of Distractors*, where Frogger is used to study mask-guided latent action learning when moving objects also drive the player's motion.

It is an independently generated research dataset and is not an official Atari, Arcade Learning Environment, OCAtari, or CleanRL release.

Preview

The preview shows a training rollout with the RGB observation, binary player mask, and mask overlay side by side.

<video controls autoplay loop muted playsinline src="https://huggingface.co/datasets/EpicPinkPenguin/atarifroggerwith_masks/resolve/main/assets/preview.mp4"></video>

Dataset usage

Streaming is recommended for inspection and for pipelines that do not require a local copy of all 2.2 million transitions:

from datasets import load_dataset

dataset = load_dataset(
    "EpicPinkPenguin/atari_frogger_with_masks",
    name="FroggerNoFrameskip-v4",
    split="train",
    streaming=True,
)

sample = next(iter(dataset))
print(sample.keys())
# dict_keys(['observation', 'mask', 'action', 'reward',
#            'terminated', 'truncated'])

print(sample["observation"].size)  # (128, 128)
print(sample["mask"].size)         # (128, 128)

Remove streaming=True to download and cache a complete split locally:

from datasets import load_dataset

test_dataset = load_dataset(
    "EpicPinkPenguin/atari_frogger_with_masks",
    name="FroggerNoFrameskip-v4",
    split="test",
)

Rows are stored in temporal order. Episodes can be reconstructed from the terminal flags:

def episodes(dataset):
    episode = []
    for step in dataset:
        episode.append(step)
        if step["terminated"] or step["truncated"]:
            yield episode
            episode = []

Dataset structure

Each row contains one observation, its player mask, the action selected from that observation, and arrival metadata for that observation.

The stored 128 × 128 RGB image is not the direct policy input. The policy acts on a stack of the four most recent 84 × 84 grayscale frames. Stored masks and RGB observations are spatially aligned at 128 × 128.

Splits

Dataset creation

The dataset was collected from a CleanRL Rainbow expert in FroggerNoFrameskip-v4, without action repetition or frame skipping. The policy consumes four stacked 84 × 84 grayscale frames, while the dataset stores 128 × 128 RGB observations. Binary player masks are programmatically derived from OCAtari vision-based object localization and player sprite color matching, then aligned with the stored observations.

Intended uses and limitations

Suitable uses include offline visual control, behavior cloning, latent-action learning, foreground-aware representation learning, segmentation evaluation, and controlled comparisons of loss masking and input masking. The pairing of the complete scene with a player-only mask is particularly useful for studying whether a representation retains interactions between the player and moving environment objects.

Important limitations:

• —The data covers the state-action distribution of a single high-performing deterministic expert, not arbitrary policies or the full game state space.

License and attribution

No dataset-wide license is declared in this repository. Users are responsible for reviewing the licenses and usage terms of CleanRL, OCAtari, the Arcade Learning Environment, and the Atari ROMs required to run the environment. Availability of the dataset does not grant rights to Atari game assets or ROMs.

Citation

Please cite the accompanying paper and the upstream work used to train the expert and produce the dataset:

@misc{fechner2026segmentfocusguidinglatent,
  title         = {Segment to Focus: Guiding Latent Action Models in the Presence of Distractors},
  author        = {Marcus Fechner and Hamza Adnan and Constantin C. L{\"u}th and Matthew T. Jackson and Alexey Zakharov and J. Marius Z{\"o}llner},
  year          = {2026},
  eprint        = {2602.02259},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2602.02259}
}

@inproceedings{hessel2018rainbow,
  title     = {Rainbow: Combining Improvements in Deep Reinforcement Learning},
  author    = {Hessel, Matteo and Modayil, Joseph and van Hasselt, Hado and Schaul, Tom and Ostrovski, Georg and Dabney, Will and Horgan, Dan and Piot, Bilal and Azar, Mohammad and Silver, David},
  booktitle = {Proceedings of the AAAI Conference on Artificial Intelligence},
  volume    = {32},
  number    = {1},
  year      = {2018}
}

@article{bellemare2013arcade,
  title   = {The Arcade Learning Environment: An Evaluation Platform for General Agents},
  author  = {Bellemare, Marc G. and Naddaf, Yavar and Veness, Joel and Bowling, Michael},
  journal = {Journal of Artificial Intelligence Research},
  volume  = {47},
  pages   = {253--279},
  year    = {2013}
}

@article{delfosse2024ocatari,
  title   = {{OCAtari}: {O}bject-Centric {Atari} 2600 Reinforcement Learning Environments},
  author  = {Delfosse, Quentin and Bl{\"u}ml, Jannis and Gregori, Bjarne and Sztwiertnia, Sebastian and Kersting, Kristian},
  journal = {Reinforcement Learning Journal},
  volume  = {1},
  pages   = {400--449},
  year    = {2024}
}

@article{huang2022cleanrl,
  title   = {{CleanRL}: High-quality Single-file Implementations of Deep Reinforcement Learning Algorithms},
  author  = {Huang, Shengyi and Dossa, Rousslan Fernand Julien and Ye, Chang and Braga, Jeff and Chakraborty, Dipam and Mehta, Kinal and Ara{\'u}jo, Jo{\~a}o G. M.},
  journal = {Journal of Machine Learning Research},
  volume  = {23},
  number  = {274},
  pages   = {1--18},
  year    = {2022}
}