p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] pyproject.toml (plaintext)
[project]
name = "jasmine"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
  "dm-pix>=0.4.3",
  "einops>=0.8.0",
  "flax>=0.10.7",
  "jax[cuda12]>=0.6.2",
  "optax>=0.2.3",
  "tyro>=0.8.5",
  "wandb>=0.17.4",
  "grain>=0.2.10",
  "array-record>=0.7.2",
  "pre-commit>=4.2.0",
]

[build-system]
requires = ["uv-build>=0.8.21,<0.9.0"]
build-backend = "uv_build"
[tab] data/pyproject.toml (plaintext)
[project]
name = "jasmine-data"
version = "0.1.0"
requires-python = "==3.10"
dependencies = [
  "array-record>=0.7.2",
  "numpy",
  "procgen>=0.10.7",
  "hf-transfer==0.1.9",
  "huggingface-hub[cli]>=0.34.3",
  "ffmpeg-python==0.2.0",
  "tqdm>=4.67.1",
]

[build-system]
requires = ["uv-build>=0.8.21,<0.9.0"]
build-backend = "uv_build"
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export PYTHONUNBUFFERED=1
srun python train_tokenizer.py \
    --save_ckpt \
    --log_checkpoint_interval 10 \
    --image_height=64 \
    --image_width=64 \
    --patch_size=16 \
    --ckpt_dir $ckpt_dir \
    --batch_size 120 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --data_dir $data_dir
[tab] slurm/dev/mihir/horeka/train_dynamics.sh (shellscript)
#!/usr/bin/env bash

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


tf_records_dir=$ws_dir/knoms_tfrecords_500_shards
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'


job_name="debug"
slurm_job_id="0000"

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] README.md (markdown)
bash input_pipeline/download/huggingface/download_openai_array_records.sh
```

---

### Option 2: Manual Download & Preprocessing of OpenAI's VPT Dataset

If you prefer to use the raw VPT dataset from OpenAI and preprocess it yourself, follow these steps:

1. **Download index files:**
   This will download the initial index file:

   ```bash
   bash input_pipeline/download/openai/download_index_files.sh
   ```

2. **Download from all index files:**
   This may take a long time depending on your bandwidth:

   ```bash
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_7xx_Apr_6.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_8xx_Jun_29.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_9xx_Jun_29.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_10xx_Jun_29.json
   ```

3. **Preprocess videos into ArrayRecords:**
   For efficient distributed training, convert the raw videos into the arrayrecord format (make sure to have [ffmpeg](https://github.com/FFmpeg/FFmpeg) installed on your machine):

   ```bash
   python input_pipeline/preprocess/video_to_array_records.py
[tab] pyproject.toml (plaintext)
[project]
name = "jasmine"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
  "dm-pix>=0.4.3",
  "einops>=0.8.0",
  "flax>=0.10.7",
  "jax[cuda12]>=0.6.2",
  "optax>=0.2.3",
  "tyro>=0.8.5",
  "wandb>=0.17.4",
  "grain>=0.2.10",
  "array-record>=0.7.2",
  "pre-commit>=4.2.0",
  "pillow>=11.3.0",
]
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
ckpt_dir="${PWD}/checkpoints/dynamics_openai_grain_tok_restore"
tokenizer_ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

export PYTHONUNBUFFERED=1
srun python jasmine/train_dynamics.py \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --dyna_type 'causal' \
    --batch_size 120 \
    --image_height=64 \
    --image_width=64 \
    --patch_size=16 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] data/pyproject.toml (plaintext)
[project]
name = "jasmine-data"
version = "0.1.0"
requires-python = "==3.10"
dependencies = [
  "array-record>=0.7.2",
  "numpy",
  "procgen>=0.10.7",
  "hf-transfer==0.1.9",
  "huggingface-hub[cli]>=0.34.3",
  "ffmpeg-python==0.2.0",
  "pillow>=11.3.0",
  "tqdm>=4.67.1",
]
[tab] data/utils.py (python)
import os
import pickle
import numpy as np
from array_record.python.array_record_module import ArrayRecordWriter


def save_chunks(file_idx, chunks_per_file, output_dir, obs_chunks, act_chunks=None):
    os.makedirs(output_dir, exist_ok=True)

    metadata = []
    while len(obs_chunks) >= chunks_per_file:
        chunk_batch = obs_chunks[:chunks_per_file]
        obs_chunks = obs_chunks[chunks_per_file:]
        act_chunk_batch = None
        if act_chunks:
            act_chunk_batch = act_chunks[:chunks_per_file]
            act_chunks = act_chunks[chunks_per_file:]
        episode_path = os.path.join(output_dir, f"data_{file_idx:04d}.array_record")
        writer = ArrayRecordWriter(str(episode_path), "group_size:1")
        seq_lens = []
        for idx, chunk in enumerate(chunk_batch):
            seq_len = chunk.shape[0]
            seq_lens.append(seq_len)
            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": seq_len,
            }
            if act_chunk_batch:
                assert len(chunk) == len(
                    act_chunk_batch[idx]
                ), f"Observation data length and action sequence length do not match: {len(chunk)} != {len(act_chunk_batch[idx])}"
                chunk_record["actions"] = act_chunk_batch[idx]
            writer.write(pickle.dumps(chunk_record))
        writer.close()
        file_idx += 1
        metadata.append(
            {
                "path": episode_path,
                "num_chunks": len(chunk_batch),
                "avg_seq_len": np.mean(seq_lens),
            }
        )
        print(f"Created {episode_path} with {len(chunk_batch)} video chunks")

    return metadata, file_idx, obs_chunks, act_chunks
[tab] data/jasmine_data/minecraft/huggingface/download_openai_array_records.sh (shellscript)
#!/bin/bash

# Download and extract array records from Hugging Face
# 
# This script performs a two-step process:
# 1. Downloads compressed array records from a Hugging Face dataset repository
# 2. Extracts the compressed tar files in parallel for better performance
#
# Usage:
#   ./download_openai_array_records.sh [hf_download_dir] [final_dataset_dir]
#
# Arguments:
#   hf_download_dir    - Directory to store compressed downloads (default: data/minecraft_arrayrecords_compressed)
#   final_dataset_dir  - Directory for extracted array records (default: data/minecraft_arrayrecords)

# Set default directories if not provided as arguments
hf_download_dir="${1:-data/minecraft_arrayrecords_compressed}" 
final_dataset_dir="${2:-data/minecraft_arrayrecords}"          

mkdir -p $hf_download_dir
mkdir -p $final_dataset_dir

# Step 1: Download compressed dataset from Hugging Face
echo "Starting download from Hugging Face..."
repo_id=p-doom/open_ai_minecraft_arrayrecords_chunked
start_time_hf_download=$(date +%s)

HF_HUB_ENABLE_HF_TRANSFER=1 HF_HUB_DISABLE_SYMLINKS=1 \
huggingface-cli download --repo-type dataset $repo_id --local-dir $hf_download_dir

end_time_hf_download=$(date +%s)
echo "Download completed. Time taken: $((end_time_hf_download - start_time_hf_download)) seconds"

# Step 2: Extract compressed array records in parallel
echo "Starting parallel extraction of tar files..."
num_workers=64  # Number of parallel extraction processes
start_time_uncompress=$(date +%s)

# Find all shard tar files and extract them in parallel:
xargs -0 -P $num_workers -I {} bash -c 'echo "Extracting {}"; tar -xf "{}" -C "'$final_dataset_dir'"'

end_time_uncompress=$(date +%s)

# Display timing summary
echo "================================"
echo "Extraction completed successfully!"
echo "Uncompress time: $((end_time_uncompress - start_time_uncompress)) seconds"
echo "Download time: $((end_time_hf_download - start_time_hf_download)) seconds"
echo "Total time: $((end_time_uncompress - start_time_hf_download)) seconds"
echo "Final dataset location: $final_dataset_dir"
[tab] data/jasmine_data/minecraft/openai/download_index_files.sh (shellscript)
#!/bin/bash
# Download index files from OpenAI Video-Pre-Training dataset
# https://github.com/openai/Video-Pre-Training
#
# This script downloads specific JSON index files containing metadata
# for the OpenAI Video-Pre-Training dataset from Microsoft Azure blob storage.
#
# Usage:
#   ./download_index_files.sh [output_dir]
#
# Arguments:
#   output_dir: Directory to save downloaded files (default: data/open_ai_index_files)
#
# Example:
#   ./download_index_files.sh /path/to/custom/directory

# Set output directory, use default if not provided
output_dir="${1:-data/open_ai_index_files}" 
mkdir -p $output_dir

# List of index files to download
# These files contain metadata for different video ranges (6xx, 7xx, etc.)
index_files=(
    "all_6xx_Jun_29.json"    
    "all_7xx_Apr_6.json"     
    "all_8xx_Jun_29.json"    
    "all_9xx_Jun_29.json"    
    "all_10xx_Jun_29.json"   
)

# Download each index file from Azure blob storage
for index_file in "${index_files[@]}"; do
    echo "Downloading $index_file..."
    wget https://openaipublic.blob.core.windows.net/minecraft-rl/snapshots/$index_file -O $output_dir/$index_file
done

echo "Download complete. Files saved to: $output_dir"
[content] README.md (markdown)
Jasmine requires `python 3.10`, `jax 0.6.2`, and `flax 0.10.7`. To install the requirements, run:

```bash
pip install -r requirements.txt
pre-commit install
```

---

<h2 name="dataset" id="dataset">Dataset 📂</h2>

You can either download our preprocessed dataset from [Hugging Face](https://huggingface.co/datasets/p-doom/open_ai_minecraft_arrayrecords_chunked) or preprocess [OpenAI's VPT dataset](https://github.com/openai/Video-Pre-Training) manually.

### Option 1: Use Preprocessed Dataset (Recommended)

The easiest way to get started is to download our preprocessed dataset from Hugging Face. This script will handle downloading and extracting it:

```bash
bash input_pipeline/download/huggingface/download_openai_array_records.sh
```

---

### Option 2: Manual Download & Preprocessing of OpenAI's VPT Dataset

If you prefer to use the raw VPT dataset from OpenAI and preprocess it yourself, follow these steps:

1. **Download index files:**
   This will download the initial index file:

   ```bash
   bash input_pipeline/download/openai/download_index_files.sh
   ```

2. **Download from all index files:**
   This may take a long time depending on your bandwidth:

   ```bash
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_7xx_Apr_6.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_8xx_Jun_29.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_9xx_Jun_29.json
   python input_pipeline/download/openai/download_videos.py --index_file_path data/open_ai_index_files/all_10xx_Jun_29.json
   ```

3. **Preprocess videos into ArrayRecords:**
   For efficient distributed training, convert the raw videos into the arrayrecord format (make sure to have [ffmpeg](https://github.com/FFmpeg/FFmpeg) installed on your machine):

   ```bash
   python input_pipeline/preprocess/video_to_array_records.py
   ```

> **Note:** This is a large dataset and may take considerable time and storage to download and process.


<h2 name="train" id="train">Quick Start 🚀 </h2>

Genie has three components: a [video tokenizer](models/tokenizer.py), a [latent action model](models/lam.py), and a [dynamics model](models/dynamics.py). Each of these components are trained separately, however, the dynamics model requires a pre-trained video tokenizer (and latent action model).

To train the video tokenizer, run:

```bash
python train_tokenizer.py --ckpt_dir <path>
```

To train the latent action model, run:

```bash
python train_lam.py --ckpt_dir <path>
```

Once the tokenizer and LAM are trained, the dynamics model can be trained with:

```bash
python train_dynamics.py --tokenizer_checkpoint <path> --lam_checkpoint <path>
```

Logging with `wandb` is supported. To enable logging, set the `WANDB_API_KEY` environment variable or run:

```bash
wandb login
```

Training can then be logged by setting the `--log` flag:

```bash
python train_tokenizer.py --log --entity <wandb-entity> --project <wandb-project>
[content] README.md (markdown)
Jasmine requires `python 3.11`, `jax 0.7.2`, and `flax 0.11.2`. To install the requirements, run:

```bash
uv sync
pre-commit install
```

---

<h2 name="dataset" id="dataset">Dataset 📂</h2>

You can either download our preprocessed dataset from [Hugging Face](https://huggingface.co/datasets/p-doom/open_ai_minecraft_arrayrecords_chunked) or preprocess [OpenAI's VPT dataset](https://github.com/openai/Video-Pre-Training) manually.

### Option 1: Use Preprocessed Dataset (Recommended)

The easiest way to get started is to download our preprocessed dataset from Hugging Face. This script will handle downloading and extracting it:

```bash
bash data/minecraft/huggingface/download_openai_array_records.sh
```

---

### Option 2: Manual Download & Preprocessing of OpenAI's VPT Dataset

If you prefer to use the raw VPT dataset from OpenAI and preprocess it yourself, follow these steps:

1. **Download index files:**
   This will download the initial index file:

   ```bash
   bash data/minecraft/openai/download_index_files.sh
   ```

2. **Download from all index files:**
   This may take a long time depending on your bandwidth:

   ```bash
   cd data/
   uv run python jasmine_data/minecraft/openai/download_videos.py --index_file_path data/open_ai_index_files/all_7xx_Apr_6.json
   uv run python jasmine_data/minecraft/openai/download_videos.py --index_file_path data/open_ai_index_files/all_8xx_Jun_29.json
   uv run python jasmine_data/minecraft/openai/download_videos.py --index_file_path data/open_ai_index_files/all_9xx_Jun_29.json
   uv run python jasmine_data/minecraft/openai/download_videos.py --index_file_path data/open_ai_index_files/all_10xx_Jun_29.json
   ```

3. **Preprocess videos into ArrayRecords:**
   For efficient distributed training, convert the raw videos into the arrayrecord format (make sure to have [ffmpeg](https://github.com/FFmpeg/FFmpeg) installed on your machine):

   ```bash
   cd data/
   uv run python jasmine_data/video_to_array_records.py
   ```

> **Note:** This is a large dataset and may take considerable time and storage to download and process.


<h2 name="train" id="train">Quick Start 🚀 </h2>

Genie has three components: a [video tokenizer](models/tokenizer.py), a [latent action model](models/lam.py), and a [dynamics model](models/dynamics.py). Each of these components are trained separately, however, the dynamics model requires a pre-trained video tokenizer (and latent action model).

To train the video tokenizer, run:

```bash
uv run python jasmine/train_tokenizer.py --ckpt_dir <path>
```

To train the latent action model, run:

```bash
uv run python jasmine/train_lam.py --ckpt_dir <path>
```

Once the tokenizer and LAM are trained, the dynamics model can be trained with:

```bash
uv run python jasmine/train_dynamics.py --tokenizer_checkpoint <path> --lam_checkpoint <path>
```

Logging with `wandb` is supported. To enable logging, set the `WANDB_API_KEY` environment variable or run:

```bash
wandb login
```

Training can then be logged by setting the `--log` flag:

```bash
uv run python jasmine/train_tokenizer.py --log --entity <wandb-entity> --project <wandb-project>
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export PYTHONUNBUFFERED=1
srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --restore_ckpt \
    --log_checkpoint_interval 10 \
    --image_height=64 \
    --image_width=64 \
    --patch_size=16 \
    --ckpt_dir $ckpt_dir \
    --batch_size 120 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --data_dir $data_dir
[tab] slurm/dev/alfred/berlin/gt_actions/train_dynacmis_on_3nodes_2gpu.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=3
#SBATCH --ntasks-per-node=1
#SBATCH --time=0:30:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:2
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/ali/minecraft/dynamics/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/ali/minecraft/dynamics/%x_%j.log
#SBATCH --job-name=minecraft_dynamics_on_2nodes
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout



########################################################
# setup the environment
########################################################
cat $0

source .venv/bin/activate

job_name="$SLURM_JOB_NAME"
tags="dynamics minecraft topology_restore 8gpu"

# array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/avocadoali/overfit_set"
array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/avocadoali/test_set_uncompressed"
# tokenizer_checkpoint_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/minecraft/tokenizer/overfit/interactive_20640"
tokenizer_checkpoint_dir=/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/minecraft/tokenizer/overfit/minecraft_tokenizer_topology_test_two_gpus_20687
# checkpoint_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/dev/${job_name}_${slurm_job_id}"
# mkdir -p $checkpoint_dir

checkpoint_dir=/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/dev/minecraft_dynamics_single_gpu_overfit_/000440_ckpt

export xla_flags=--xla_gpu_autotune_level=0

# srun ipython --pdb train_dynamics_get_sample.py --  \
srun python jasmine/train_dynamics_get_sample.py \
    --save_ckpt \
    --restore_ckpt \
    --ckpt_dir $checkpoint_dir \
    --tokenizer_checkpoint $tokenizer_checkpoint_dir \
    --batch_size=48 \
    --log_checkpoint_interval=5 \
    --log \
    --name=${job_name} \
    --tags $tags \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[tab] pyproject.toml (plaintext)
[project]
name = "jasmine"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
  "dm-pix>=0.4.3",
  "einops>=0.8.0",
  "flax>=0.10.7",
  "jax[cuda12]>=0.6.2",
  "optax>=0.2.3",
  "tyro>=0.8.5",
  "wandb>=0.17.4",
  "grain>=0.2.10",
  "array-record>=0.7.2",
  "pre-commit>=4.2.0",
  "pillow>=11.3.0",
]

[build-system]
requires = ["uv_build>=0.8.22,<0.9.0"]
build-backend = "uv_build"

[tool.uv.build-backend]
module-root = ""
[tab] data/pyproject.toml (plaintext)
[project]
name = "jasmine-data"
version = "0.1.0"
authors = [
  {name = "Alfred Nguyen", email = "alfrednguyen02@gmail.com"},
  {name = "Mihir Mahajan", email = "mihir@pdoom.org"},
  {name = "Franz Srambical", email = "franz@pdoom.org"},
]
requires-python = "==3.10.*"
dependencies = [
  "procgen>=0.10.7",
  "gym3>=0.3.3",
  "array-record>=0.7.2",
  "numpy",
  "hf-transfer==0.1.9",
  "huggingface-hub[cli]>=0.34.3",
  "ffmpeg-python==0.2.0",
  "pillow>=11.3.0",
  "tqdm>=4.67.1",
  "tyro>=0.8.5",
]

[build-system]
requires = ["uv_build>=0.8.22,<0.9.0"]
build-backend = "uv_build"

[tool.uv.build-backend]
module-root = ""
[content] data/jasmine_data/atari/atari_utils.py (python)
        assert (
            env.observation_space.dtype is not None
        ), "No dtype specified for the observation space"
        assert (
            env.observation_space.shape is not None
        ), "No shape defined for the observation space"
        self._obs_buffer = np.zeros(
            (2, *env.observation_space.shape), dtype=env.observation_space.dtype
        )
        self._skip = skip

    def step(self, action: int):
        """
        Step the environment with the given action
        Repeat action, sum reward, and max over last observations.

        :param action: the action
        :return: observation, reward, terminated, truncated, information
        """
        total_reward = 0.0
        terminated = truncated = False
        for i in range(self._skip):
            obs, reward, terminated, truncated, info = self.env.step(action)
            done = terminated or truncated
            if i == self._skip - 2:
                self._obs_buffer[0] = obs
            if i == self._skip - 1:
                self._obs_buffer[1] = obs
            total_reward += float(reward)
            if done:
                break
        # Note that the observation on the done=True frame
        # doesn't matter
        max_frame = self._obs_buffer.max(axis=0)

        return max_frame, total_reward, terminated, truncated, info


class ClipRewardEnv(gym.RewardWrapper):
    """
    Clip the reward to {+1, 0, -1} by its sign.

    :param env: Environment to wrap
    """

    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)

    def reward(self, reward: SupportsFloat) -> float:
        """
        Bin reward to {+1, 0, -1} by its sign.

        :param reward:
        :return:
        """
        return np.sign(float(reward))


class WarpFrame(gym.ObservationWrapper[np.ndarray, int, np.ndarray]):
    """
    Convert to grayscale and warp frames to 84x84 (default)
    as done in the Nature paper and later work.

    :param env: Environment to wrap
    :param width: New frame width
    :param height: New frame height
    """

    def __init__(self, env: gym.Env, width: int = 84, height: int = 84) -> None:
        super().__init__(env)
        self.width = width
        self.height = height
        assert isinstance(
            env.observation_space, spaces.Box
        ), f"Expected Box space, got {env.observation_space}"

        self.observation_space = spaces.Box(
            low=0,
            high=255,
            shape=(self.height, self.width, 1),
            dtype=env.observation_space.dtype,  # type: ignore[arg-type]
        )

    def observation(self, frame: np.ndarray) -> np.ndarray:
        """
        returns the current observation from a frame

        :param frame: environment frame
        :return: the observation
        """
        assert (
            cv2 is not None
        ), "OpenCV is not installed, you can do `pip install opencv-python`"
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        frame = cv2.resize(
            frame, (self.width, self.height), interpolation=cv2.INTER_AREA
        )