recap-value-network-so101-pickplace

RECAP Value Network

A distributional value network trained with the RECAP (Reinforcement Learning from Corrective Actions and Preferences) framework on the jackvial/so101_pickplace_recap_merged_v2 dataset.

This model predicts per-frame normalized expected returns for robot manipulation trajectories, discretized into 50 bins over [-1.0, 0.0].

Model Details

Training Metrics (Epoch 1)

Usage

Loading the model

import json
import torch
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

from lerobot.rl.algorithms.recap_value_network import (
    RECAPValueNetwork,
    RECAPValueNetworkConfig,
)

# Download and load config
config_path = hf_hub_download("jackvial/recap-value-network-so101-pickplace", "config.json")
with open(config_path) as f:
    config_dict = json.load(f)

# Reconstruct the config (filter to only RECAPValueNetworkConfig fields)
from dataclasses import fields as dc_fields
valid_keys = {f.name for f in dc_fields(RECAPValueNetworkConfig)}
config = RECAPValueNetworkConfig(
    **{k: v for k, v in config_dict.items() if k in valid_keys}
)
# Don't re-download pretrained VLM weights during loading
config.pretrained_path = None

# Build model and load weights
model = RECAPValueNetwork(config)
weights_path = hf_hub_download("jackvial/recap-value-network-so101-pickplace", "model.safetensors")
state_dict = load_file(weights_path)
model.load_state_dict(state_dict)
model.eval()

Using as a value network for RECAP pi-star training

uv run python -m lerobot.rl.algorithms.recap_train_pi_star \
  --repo_id="jackvial/so101_pickplace_recap_merged_v2" \
  --output_dir="outputs/recap_pistar" \
  --value_network_checkpoint="jackvial/recap-value-network-so101-pickplace"

Training Configuration

Trained with lerobot.rl.algorithms.recap_train_value_network:

{
  "repo_id": "jackvial/so101_pickplace_recap_merged_v2",
  "output_dir": "/home/jack/code/lerobot/outputs/recap_value_5",
  "labels_csv_path": null,
  "root": "/home/jack/.cache/huggingface/lerobot",
  "revision": null,
  "episodes": null,
  "epochs": 2,
  "batch_size": 4,
  "gradient_accumulation_steps": 4,
  "num_workers": 4,
  "learning_rate": 0.0001,
  "weight_decay": 0.0001,
  "warmup_ratio": 0.05,
  "max_grad_norm": 1.0,
  "val_split_ratio": 0.1,
  "seed": 42,
  "device": "auto",
  "max_train_steps_per_epoch": null,
  "max_val_steps_per_epoch": null,
  "log_every_n_steps": 100,
  "validate_every_n_train_steps": 50,
  "plot_every_n_train_steps": 200,
  "max_val_steps_per_step_validation": 20,
  "val_plot_num_episodes": 4,
  "val_plot_num_frames": 8,
  "val_plot_every_n_epochs": 1,
  "c_fail": 500.0,
  "num_value_bins": 50,
  "tokenizer_max_length": 200,
  "image_size": 224,
  "paligemma_variant": "gemma_2b",
  "tokenizer_name": "google/paligemma-3b-pt-224",
  "model_precision": "bfloat16",
  "freeze_vision_encoder": true,
  "freeze_backbone": true,
  "num_unfrozen_backbone_layers": 3,
  "num_vlm_layers": 10,
  "value_head_depth": 1,
  "dropout": 0.1,
  "pretrained_path": "lerobot/pi05_base",
  "wandb_project": null,
  "wandb_entity": null,
  "wandb_run_name": null
}