Skip to content

About

Official implementation of the paper "Rank-Then-Act: Reward-Free Control from Frame-Order Progress"

Topics

Resources

Stars

6 stars

Watchers

0 watching

Forks

Latest commit

 

History

7 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Rank-Then-Act

Reward-Free Control from Frame-Order Progress

Yuriy Maksyuta, George Bredis, Ruslan Rakhimov, Daniil Gavrilov

ICML 2026 Workshop on RLxF: Reinforcement Learning from World Feedback

arXiv

Rank-Then-Act is a two stage method for producing a dense rewards for downstream RL tasks. Method overview can be seen in the image below.

Method Overview

Screenshot

Our method shows great performance both in discrete environments (retro GameBoy games) and in continuous environments (Metaworld). Additionally, we have found out that RTA generalizes well to unseen tasks. Main results are presented in tables below. For more information about method, results and ablation studies please refer to the paper.

Results

Training code for vision-language models that estimate task completion from video frames and use that signal to train agents in Game Boy environments.

The repository is split into two stages:

  • Stage 1 — Fine-tune a VLM with GRPO on gameplay videos. The model learns to predict task-completion percentages (or ranks) for shuffled frame sequences. Training is guided by a VOC (rank-correlation) reward.
  • Stage 2 — Use the stage-1 model as a reward model inside RL (PPO) to train a VLM policy that plays Game Boy games (Catrap, Kirby) through a PyBoy emulator.

Requirements

  • Python 3.10+
  • CUDA-capable GPU (multi-GPU supported for stage 1)
  • uv for dependency management
  • Optional: Weights & Biases account for experiment logging

Stage 2 additionally requires:

  • xvfb for headless Game Boy rendering
  • Game Boy cartridge (.gb) and level save-state files in stage2/level_states/

Repository structure

.
├── README.md
├── .gitignore
├── stage1/                        # GRPO fine-tuning on videos
│   ├── train_grpo_v2.py           # Main training entry point
│   ├── rl_utils.py                # Templates for VLM queries
│   ├── video_datasets.py          # Video segment datasets
│   ├── video_utils.py             # Frame extraction utilities
│   ├── upload_qwen_to_s3.py       # Optional code to upload run artifacts to s3
│   ├── configs/
│   │   ├── hydra/                 # Hydra experiment configs
│   │   └── accelerate/            # Multi-GPU DeepSpeed configs
│   ├── patches/                   # Patched TRL GRPO trainer files
│   ├── scripts/run.sh             # Full setup + launch script
│   ├── splits/                    # Train / test / exclude lists for YT videos
│   ├── pyproject.toml
│   └── uv.lock
└── stage2/                        # RL with VLM-based VOC rewards
    ├── main.py                    # PPO training with VOC reward model
    ├── pyboy_gym.py               # PyBoy ↔ Gymnasium wrapper
    ├── voc_calculation.py         # VOC reward function used at runtime
    ├── a2c_ppo_acktr/             # PPO / A2C implementation and VLM policy
    ├── configs/                   # Hydra configs (catrap, kirby, grpo variants)
    ├── scripts/
    │   ├── run.sh                 # PPO launch script
    ├── level_states/              # Game Boy ROMs and save states (user-provided)
    ├── pyproject.toml
    └── uv.lock

Paths such as vids/catrap_level_2.mp4 in config files are relative to the stage root.

Setup

Each stage has its own environment. Install dependencies from inside the stage directory:

# Stage 1
cd stage1
uv venv --python 3.10
uv pip install --no-cache-dir torch==2.7.1
uv sync --frozen
source .venv/bin/activate
# Stage 2
cd stage2
uv venv --python 3.10
uv pip install --no-cache-dir torch==2.7.1
uv sync --frozen
source .venv/bin/activate

Stage-1 patch the TRL library with custom GRPO trainer code. After activating the environment, copy the patches:

# Run from inside stage1/ or stage2/
TRL_PATH="$(uv pip show trl | grep Location | awk '{print $2}')/trl"
cp patches/grpo_trainer.py "$TRL_PATH/trainer/grpo_trainer.py"
cp patches/grpo_config.py  "$TRL_PATH/trainer/grpo_config.py"

Running stage 1 — video GRPO training

From the stage1/ directory, launch training with Accelerate:

cd stage1
source .venv/bin/activate

export WANDB_API_KEY=<your-key>   # optional

CUDA_VISIBLE_DEVICES=0,1,2,3 \
  accelerate launch \
    --config_file=configs/accelerate/config_zero2.yaml \
    train_grpo_v2.py

Hydra configs live in stage1/configs/hydra/. Override the default config or parameters on the command line:

# Use the Catrap-specific config
accelerate launch --config_file=configs/accelerate/config_zero2.yaml \
  train_grpo_v2.py --config-name=config_v2_catrap

# Override individual settings
accelerate launch --config_file=configs/accelerate/config_zero2.yaml \
  train_grpo_v2.py \
  video.folder_path=../vids/ \
  model.model_path=Qwen/Qwen2.5-VL-7B-Instruct

Example of a config is in config_v2.yaml.

Running stage 2 — Game Boy RL with VOC rewards

Place cartridge ROMs and level save states in stage2/level_states/ before training. Configs reference paths like ./level_states/catrap.gb and ./level_states/state_file_4_f.state.

PPO (main.py)

cd stage2
source .venv/bin/activate

export WANDB_API_KEY=<your-key>   # optional

CUDA_VISIBLE_DEVICES=0 \
  xvfb-run -a -s "-screen 0 1024x768x24 -ac +extension GLX +render -noreset" \
  accelerate launch --config_file=configs/config_zero2.yaml main.py

Select a config with --config-name:

accelerate launch --config_file=configs/config_zero2.yaml main.py --config-name=vldac_catrap
accelerate launch --config_file=configs/config_zero2.yaml main.py --config-name=vldac_kirby

Set model.lora_path in the config to a stage-1 checkpoint so the reward model uses your fine-tuned weights.

About

Official implementation of the paper "Rank-Then-Act: Reward-Free Control from Frame-Order Progress"

Topics

Resources

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages