🌟 Uncertainty-Aware Policy Steering (UPS)

Paper: When to Act, Ask, or Learn: Uncertainty-Aware Policy Steering
arXiv: https://arxiv.org/abs/2602.22474
Project website: https://jessie-yuan.github.io/ups/

Authors: Jessie Yuan¹*, Yilin Wu¹*, Andrea Bajcsy¹
¹ Carnegie Mellon University

---

📌 Overview

This repository contains the official code for UPS, a framework that combines:

• 🌍 World Model (Dreamer-v3 style) trained on simulated robot rollouts & demonstrations to imagine future outcomes of action sequences as videos
• ♊ VLM Narrator + Calibrated Verifier (Gemini 3 flash preview, out-of-the-box) that consumes decoded WM imaginations, narrates them into text, and then selects a set of narrations such that a narration corresponding to an action sequence that appropriately responds to the user's intent is containeed within this set with high probability. This verfication portion must be calibrated through conformal prediction.

This repository enables you to explore this framework for the RoboMimic NutAssembly task in simulation. We use the same pipeline for our hardware experiments, but for simplicity this repository only looks at how to run UPS in simulation.

---

🗂️ Repository Structure

• diffusion4robotics: Code needed to train the base diffusion policy, the intervention classifer, and the residual policy.
• model_based_irl_torch: Code needed to train the world model.
• policy: Wrapper classes for the base / residual policies, the WM predictor, the VLM translator, and the conformal-prediction based VLM verifier.
• scripts: Scripts to run the conformal prediction calibration step, correction data processing, and the overall pipeline.
• vlm_query: Prompts and reference images.

---

⬇️ Installation

1) Clone the repo (with submodules)

git clone --recurse-submodules git@github.com:CMU-IntentLab/uncertainty_aware_policy_steering.git
cd ups

2) Create the conda environment

create env + pip install

conda env create -f environment.yaml
conda activate ups
pip install -r requirements.txt

install dreamer wm

cd model_based_irl_torch
pip install -e .

install robomimic for simulation

cd robosuite
pip install -e .

install diffusion policy

cd diffusion4robotics
pip install git+https://github.com/AGI-Labs/robobuf.git
pip install git+https://github.com/facebookresearch/r3m.git
pip install -e .

install policy

cd policy
pip install -e .

---

📦 Prepare Models

Option 1: Download the pre-trained checkpoints for the NutAssembly task

We provide pretrained checkpoints on Hugging Face: https://huggingface.co/jzyuan04/ups_nutassembly_diffusion_checkpoints for the diffusion and residual checkpoints and https://huggingface.co/jzyuan04/ups_nutassembly_wm_checkpoint for the WM checkpoints respectively. Paths to these checkpoints are passed to our scripts via flags.

After calibration, our qhat=0.7, but it is easy to run this yourself as outlined in option 2, step 3.

---

Option 2: Download our existing data for the NutAssembly task and train your own WM and diffusion policies

1. Train a diffusion policy with our demonstration data:

Download the hdf5 file containing all our demos: https://huggingface.co/jzyuan04/ups_nutassembly_base_policy_demos

cd diffusion4robotics
python combine_traj_to_buffer.py --base --hdf5-files /path/to/demo_data.hdf5 --output_dir /path/to/buffer_dir

The combine_traj_to_buffer.py script converts the hdf5 file with our demos into a buffer.pkl file, as well as json noramlization files. The path to the resultant buffer.pkl file is passed to the following script that trains the diffusion policy.

nice -n 19 python finetune.py agent=diffusion_unet buffer_path=/path/to/buffer.pkl max_iterations=100000  trainer=bc_cos_sched ac_chunk=16 train_transform=medium task.train_buffer.cam_indexes=[0,1] img_chunk=2

2. Train a WM with our WM data (demos + rollouts):

Download the hdf5 file containing all our demos + rollouts: https://huggingface.co/jzyuan04/ups_nutassembly_wm_data

Edit wm_nutassembly_config.yaml to include the path to this data.

cd model_based_irl_torch
python scripts/train_wm_real_data.py --config_path wm_nutassembly_config.yaml

3. Run the calibration step of conformal predition to determine your threshold (qhat) value:

python scripts/generate_calibration_data.py --output /path/to/calibration_narrations.json
python scripts/run_calibration --narrations /path/to/calibration_narrations.json --instructions /path/to/instructions.json

4. Collect and process correction data:

We collect and process 40 corrections using our UPS pipeline paired with a 3Dconnexion SpaceMouse (although the keyboard may be used instead). More information about how to set this up and how the controls map to the simulated robot is available here: https://robosuite.ai/docs/modules/devices.html.

python scripts/run_ups --qhat your_calibrated_qhat --corrections_dataset_path /path/to/corrections.hdf5 --wm_ckpt /path/to/your_wm.pt --base_policy_ckpt /path/to/your_base_policy.ckpt --device spacemouse (or keyboard)
python scripts/process_corrections --corrections_dataset_path /path/to/corrections.hdf5 --base_policy_ckpt /path/to/your_base_policy.ckpt

5. Train intervention classifier and residual policy:

Training these models is similar to training the diffusion policy; however, for the classifier, the weights of each class need to be computed via compute_class_weights.py, which prints a weight value to the terminal that can be passed to the training script via a flag.

cd diffusion4robotics
python combine_traj_to_buffer.py --classifier --hdf5-files /path/to/corrections.hdf5 --output_dir /path/to/classifier_buffer_dir
python compute_class_weights.py --config-name=finetune_classifier
python finetune.py --config-name=finetune_classifier agent.pos_weight=<your_weight> buffer_path=/path/to/classifier_buffer/buffer.pkl

python combine_traj_to_buffer.py --residual --hdf5-files /path/to/corrections.hdf5 --output_dir /path/to/residual_buffer_dir
python finetune.py --config-name=finetune_residual_policy buffer_path=/path/to/residual_buffer/buffer.pkl

6. Retrain WM with residual policy rollouts:

Download the hdf5 file containing our residual_policy rollouts: https://huggingface.co/jzyuan04/ups_nutassembly_wm_data_residual

Edit wm_nutassembly_config.yaml to include the path to this data.

cd model_based_irl_torch
python scripts/train_wm_real_data.py --config_path wm_nutassembly_config.yaml

---

Option 3: Create your own dataset

Follow the steps for option 2, but instead of downloading our datasets from HF, collect your own demonstrations as outlined below:

1. Collect around 100 demonstrations to train your base policy. Half should complete it with one mode and the other half should complete it with another (e.g. handle facing left vs handle facing right).
2. Collect around 250 rollouts of this base policy. Ideally both modes demonstrated and failures should be represented.

Follow steps 3-5 in option 2.

1. Collect 100 rollouts of the residual policy and retrain the WM with the original demonstration data, the base policy rollouts, and the residual policy rollouts combined.

---

🏃 Run the overall pipeline:

python scripts/run_ups --qhat your_calibrated_qhat --corrections_dataset_path /path/to/corrections.hdf5 --wm_ckpt /path/to/your_wm.pt --base_policy_ckpt /path/to/your_base_policy.ckpt --classifier_ckpt /path/to/your_classifier.ckpt --residual_ckpt /path/to/your_residual_policy.ckpt --device spacemouse (or keyboard)

---

📚 Citation

If you find this work useful, please cite:

@article{yuan2026ups,
    title={When to Act, Ask, or Learn: Uncertainty-Aware Policy Steering},
    author={Yuan, Jessie and Wu, Yilin and Bajcsy, Andrea},
    journal={Robotics: Science and Systems (RSS)},
    year={2026}
}