gym_wm - World Model Dataset Generator

A Python toolbox for generating offline datasets from Gymnasium environments for world-model research using Minari.

Features

• 🎮 Multi-environment support: Gymnasium MuJoCo, Gymnasium-Robotics (Fetch, Maze, Hand), and classic control
• 📸 RGB observations: Automatic capture of visual observations alongside proprioceptive state
• 🗄️ Minari integration: Store datasets in Minari format for easy sharing and loading
• 🎯 Multiple policies: Random sampling or PD controller for goal-conditioned environments
• 📊 Visualization tools: Episode frames, trajectory plots, and dataset summaries
• 🎬 Video export: Create MP4 videos from collected episodes
• 💻 CLI & Python API: Use from command line or import as a package

Installation

# Clone the repository
git clone https://github.com/unknown/gym_wm_dataset_generator.git
cd gym_wm_dataset_generator

# Create environment (using micromamba)
mm create -n gym_wm python=3.11
mm activate gym_wm

# Install with uv (recommended)
uv pip install -e ".[dev]"

# Or install dependencies only
uv pip install -e .

Dependencies

Core dependencies:

• gymnasium>=1.0.0
• gymnasium-robotics>=1.3.0
• minari>=0.5.0
• numpy>=1.24.0
• typer>=0.12.0
• rich>=13.0.0
• loguru>=0.7.0
• matplotlib>=3.7.0
• imageio>=2.31.0

Quick Start

Command Line Interface

# List supported environments
gym-wm list-envs

# Collect a dataset
gym-wm collect PointMaze_UMaze-v3 --episodes 100

# Collect with custom options
gym-wm collect FetchReach-v4 -n 50 --policy pd --img-size 128 --seed 42

# List collected datasets
gym-wm list-datasets

# Inspect a dataset
gym-wm inspect pointmaze-umaze-v3/random-v0

# Visualize a dataset
gym-wm visualize pointmaze-umaze-v3/random-v0 --episode 0

# Generate dataset summary
gym-wm visualize pointmaze-umaze-v3/random-v0 --summary

# Create video from episode
gym-wm visualize pointmaze-umaze-v3/random-v0 --save-video

# Compare multiple datasets
gym-wm compare dataset1 dataset2

Python API

import gym_wm

# List available environments
envs = gym_wm.list_environments()
print(envs.keys())  # ['mujoco', 'fetch', 'maze', 'hand', 'classic_control']

# Collect a dataset with random policy
dataset = gym_wm.collect_dataset(
    env_id="PointMaze_UMaze-v3",
    num_episodes=100,
    dataset_name="pointmaze/my-dataset-v0",
    img_height=84,
    img_width=84,
)
print(f"Collected {dataset.total_episodes} episodes, {dataset.total_steps} steps")

# Collect with PD controller policy (for goal-conditioned envs)
policy = gym_wm.create_pd_policy(kp=2.0, kd=0.2)
dataset = gym_wm.collect_dataset(
    env_id="FetchReach-v4",
    num_episodes=50,
    policy=policy,
)

# Load an existing dataset
dataset = gym_wm.load_dataset("pointmaze/my-dataset-v0")

# Visualize dataset
gym_wm.visualize_dataset(dataset, episode_idx=0)

# Plot specific visualizations
fig = gym_wm.plot_episode_frames(dataset, episode_idx=0)
fig = gym_wm.plot_episode_trajectory(dataset, episode_idx=0)
fig = gym_wm.plot_dataset_summary(dataset)

# Create video
video_path = gym_wm.create_episode_video(dataset, episode_idx=0, fps=30)

# Inspect dataset
info = gym_wm.inspect_dataset("pointmaze/my-dataset-v0")
print(f"Has images: {info['has_images']}")
print(f"Image shape: {info.get('image_shape')}")

Using DatasetConfig

from gym_wm import DatasetConfig, collect_dataset_from_config

config = DatasetConfig(
    env_id="Ant-v5",
    num_episodes=100,
    dataset_name="ant/random-v0",
    img_height=128,
    img_width=128,
    author="Your Name",
    author_email="your@email.com",
    seed=42,
)

dataset = collect_dataset_from_config(config)

Supported Environments

MuJoCo

• Reacher-v5, Ant-v5, HalfCheetah-v5, Hopper-v5
• Walker2d-v5, Humanoid-v5, Swimmer-v5, Pusher-v5
• InvertedPendulum-v5, InvertedDoublePendulum-v5

Fetch (Gymnasium-Robotics)

• FetchReach-v4, FetchPush-v3, FetchSlide-v3, FetchPickAndPlace-v3

Maze (Gymnasium-Robotics)

• PointMaze_UMaze-v3, PointMaze_Medium-v3, PointMaze_Large-v3
• AntMaze_UMaze-v5, AntMaze_Medium-v5, AntMaze_Large-v5

Hand (Gymnasium-Robotics)

• AdroitHandDoor-v1, AdroitHandHammer-v1
• AdroitHandPen-v1, AdroitHandRelocate-v1

Classic Control

• CartPole-v1, Pendulum-v1, Acrobot-v1
• MountainCar-v0, MountainCarContinuous-v0

CLI Reference

Common Options

--episodes, -n    Number of episodes to collect (default: 100)
--name            Custom dataset name
--img-size        Image size in pixels (default: 84)
--output, -o      Output directory
--seed            Random seed for reproducibility
--policy          Policy type: random, pd (default: random)
--verbose, -v     Enable verbose logging

Development

# Install development dependencies
uv pip install -e ".[dev,test]"

# Run tests
pytest

# Run linter
ruff check src/

# Format code
ruff format src/

# Type checking
mypy src/gym_wm

Project Structure

gym_wm_dataset_generator/
├── src/
│   └── gym_wm/
│       ├── __init__.py       # Public API exports
│       ├── cli.py            # Typer CLI application
│       ├── core/
│       │   ├── config.py     # DatasetConfig dataclass
│       │   ├── environments.py # Environment utilities
│       │   ├── generator.py  # Dataset collection
│       │   ├── policies.py   # Policy implementations
│       │   └── visualize.py  # Visualization tools
│       └── utils/
│           └── inspect.py    # Dataset inspection
├── tests/
├── pyproject.toml
└── README.md

License

MIT License - see LICENSE for details.

References

• Gymnasium Documentation
• Gymnasium-Robotics Documentation
• Minari Documentation