PAVXploreRL

Official implementation of:

PAVXploreRL: World Model Reinforcement Learning for Physical-Action-Visual Targets with Out-of-Distribution Action Exploration

Overview

PAVXploreRL is a reinforcement learning framework for training action-conditioned world models using rewards that jointly optimize:

• Physical Plausibility – generated motions remain physically consistent.
• Action Adherence – generated videos follow the input action sequence.
• Visual Fidelity – generated observations remain visually realistic.

To improve exploration during RL, PAVXploreRL leverages both:

• In-distribution actions
• Out-of-distribution (OOD) action exploration

---

Installation

Create Environment

Create and activate the conda environment:

conda env create -f env.yml -n PAVXploreRL
conda activate PAVXploreRL

Download Pretrained Models

Download all baseline models and released checkpoints:

python ./pretrained_models/baseline_download.py

mkdir -p ./pretrained_models/SAM

wget -c \
https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth \
-O ./pretrained_models/SAM/sam_vit_h_4b8939.pth

Download Datasets

We use the following public datasets:

• Hugging Face: agibot-world/AgiBotWorld-Alpha
• Hugging Face: cadene/droid

After downloading, modify the following fields in the configuration files:

dataset_dirs:
metadata_path:

to point to your local dataset locations.

---

Inference

We provide demonstration examples in Hugging Face: Wehere/PAVXploreRL_Agibot_Subset please download and put inside:

./datasets/Agibot_subset

The examples follow the same format as AgiBotWorld-Alpha.

To perform inference using our RL checkpoint with a prediction horizon of 9.6 seconds:

python ./inference/inference_videos.py \
    --noisy_action False \
    --chunk_size 9.6 \
    --config_path ./configs/agibot_rl.yaml \
    --reward_config_path ./vjepa2/configs/train/vitg16/agibot-256px-8f.yaml \
    --ckpt_dir ./pretrained_models/PAVXploreRL_RL/agibot_rl/model.safetensors \
    --output_dir_name agibot_rl

---

Evaluation

Standard Video Quality Metrics

Evaluate generated videos using:

• PSNR
• SSIM
• LPIPS
• EPE Mean
• COS Mean
• VJEPA Encoder Similarity

python ./inference/evaluate_inference_videos.py

Parameters can be adjusted directly inside the script.

VJEPA Predictor Similarity Evaluation

To evaluate VJEPA predictor rewards, noisy actions must be enabled:

python ./inference/inference_videos.py \
    --noisy_action True \
    --chunk_size 9.6 \
    --config_path ./configs/agibot_rl.yaml \
    --reward_config_path ./vjepa2/configs/train/vitg16/agibot-256px-8f.yaml \
    --ckpt_dir ./pretrained_models/PAVXploreRL_RL/agibot_rl/model.safetensors \
    --output_dir_name agibot_rl_noise

---

Reward Coefficient Selection

To reduce the computational cost of RL training, we perform reward coefficient selection using an automated classifier-based procedure.

Step 1: Generate Reward Statistics

Run:

nohup python ./inference/coefficient_selection_step1.py \
    --config_path ./configs/agibot.yaml \
    --reward_config_path ./vjepa2/configs/train/vitg16/agibot-256px-8f.yaml \
    --ckpt_dir ./pretrained_models/PAVXplore_Private_Pretrain/agibot_pretrain/pytorch_model/mp_rank_00_model_states.pt \
    --chunk_size 3.2 \
    --output_dir agibot_pretrain \
    > ./inference/coefficient_selection_agibot.log &

This script generates a log file containing reward statistics for 12 video categories, ranging from ground-truth videos to various synthetic corruption and noise variants.

An effective reward function should consistently assign the highest reward to the ground-truth videos.

Step 2: Search for Reward Coefficients

After generating the reward statistics, run:

python ./inference/coefficient_selection_step2.py

This script evaluates different combinations of reward coefficients and identifies promising reward formulations.

---

VJEPA2 Training

PAVXploreRL uses a customized VJEPA2 model as part of the reward function.

Train the VJEPA2 model using:

cd vjepa2

bash ./scripts/train.sh

---

Training

World Model Pretraining

Train the world model from scratch or resume from pretrained weights by modifying the resume field in the configuration file.

bash ./scripts/train.sh

Reinforcement Learning Fine-tuning

After approximately 80,000 pretraining steps, launch RL training:

bash ./scripts/train_rl.sh

---

Examples

Example generations produced by PAVXploreRL are shown below.

---

Citation

If you find this work useful, please consider citing:

@article{pavxplorerl2026,
  title={PAVXploreRL: World Model Reinforcement Learning for Physical-Action-Visual Targets with Out-of-Distribution Action Exploration},
  author={...},
  journal={...},
  year={2026}
}