Skip to content

About

No description, website, or topics provided.

Resources

Stars

29 stars

Watchers

0 watching

Forks

Repository files navigation

🔋 EnergyAction: Unimanual to Bimanual Composition with Energy-Based Models (CVPR 2026 Highlight)

Mingchen Song1 2, Xiang Deng1 3†, Jie Wei1, Dongmei Jiang2, Liqiang Nie1, Weili Guan1 3†
1Harbin Institute of Technology, Shenzhen   2PengCheng Laboratory   3Shenzhen Loop Area Institute
†Corresponding author

Paper arXiv GitHub

If you find this work useful for your research, please kindly cite our paper and star our repo.

📢 Updates

Introduction

We propose EnergyAction, a framework that compositionally transfers pre-trained unimanual manipulation policies to bimanual tasks via Energy-Based Models (EBMs). By modeling each arm's policy as an energy function and composing them through energy summation, EnergyAction generates coordinated dual-arm actions with temporal coherence and collision avoidance — requiring only minimal bimanual demonstration data.

Installation

1. Clone this repository

git clone https://github.com/xxx/EnergyAction.git
cd EnergyAction

2. Create conda environment

conda create -n energyaction python=3.10 -y
conda activate energyaction
pip install -e .

3. Install additional packages

pip install torch torchvision open_clip_torch zarr einops setproctitle

4. Install CoppeliaSim (required for evaluation)

Online evaluation requires CoppeliaSim and PyRep. Follow the PyRep installation guide to set up CoppeliaSim, then set the environment variables:

export COPPELIASIM_ROOT=/path/to/CoppeliaSim
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${COPPELIASIM_ROOT}
export QT_QPA_PLATFORM_PLUGIN_PATH=${COPPELIASIM_ROOT}

Datasets

We use the RLBench benchmark with Peract2 data format stored in Zarr for efficient loading.

Single-arm data (RLBench)

Used for Stage 1: pretraining the single-arm 3D Flow Actor backbone.

# Download single-arm RLBench training data
wget https://huggingface.co/katefgroup/3d_diffuser_actor/resolve/main/Peract_packaged.zip
unzip Peract_packaged.zip
rm Peract_packaged.zip

# Download single-arm RLBench evaluation data
wget https://huggingface.co/katefgroup/3d_flowmatch_actor/resolve/main/peract_test.zip
unzip peract_test.zip
rm peract_test.zip

# Convert to Zarr format
python data_processing/peract_to_zarr.py \
    --root <raw_single_arm_data_dir> \
    --tgt <zarr_output_dir>

Bimanual data (RLBench2)

Used for Stage 2: training the EBM Bimanual Composer.

# Download bimanual-arm RLBench data (Peract format)
python scripts/rlbench/download_peract2.py --root <data_dir>

# (Optional) Generate data from scratch
bash data_generation/generate_data.sh

# Convert raw bimanual data to Zarr format
python data_processing/peract2_to_zarr.py \
    --root <raw_bimanual_data_dir> \
    --tgt <zarr_bimanual_output_dir>

Pretrained Checkpoints

We release all pretrained checkpoints on Hugging Face: 🤗 mingchens/EnergyAction

Checkpoint Description Size
pretrain_unimanual_checkpoint.pth Stage 1: Single-arm 3D Flow Actor 421 MB
20demo_checkpoint.pth Stage 2: EBM Bimanual Composer (20 demos) 1.69 GB
100demo_checkpoint.pth Stage 2: EBM Bimanual Composer (100 demos) 1.69 GB
# Download checkpoints (requires huggingface_hub)
pip install huggingface_hub

# Download all checkpoints
huggingface-cli download mingchens/EnergyAction --local-dir checkpoints/

# Or download individually
huggingface-cli download mingchens/EnergyAction pretrain_unimanual_checkpoint.pth --local-dir checkpoints/
huggingface-cli download mingchens/EnergyAction 100demo_checkpoint.pth --local-dir checkpoints/
  • pretrain_unimanual_checkpoint.pth: Use as LEFT_PRETRAINED_PATH and RIGHT_PRETRAINED_PATH in Stage 2 training. You can skip Stage 1 entirely with this checkpoint.
  • 20demo_checkpoint.pth: EBM Bimanual Composer trained with 20 bimanual demonstrations per task.
  • 100demo_checkpoint.pth: EBM Bimanual Composer trained with 100 bimanual demonstrations per task.

Evaluation

Online Evaluation (RLBench + CoppeliaSim)

We evaluate on 13 bimanual tasks in RLBench using CoppeliaSim simulation:

# Evaluate with experiment name (auto-finds best.pth)
bash scripts/rlbench/eval_ebm_bimanual_composer.sh my_experiment

# Evaluate with specific checkpoint
bash scripts/rlbench/eval_ebm_bimanual_composer.sh /path/to/checkpoint.pth

# With nohup (recommended for long evaluation)
nohup bash scripts/rlbench/eval_ebm_bimanual_composer.sh my_experiment > eval_my_experiment.log 2>&1 &
Evaluated Bimanual Tasks (13 tasks)
Task Task
bimanual_push_box bimanual_lift_ball
bimanual_dual_push_buttons bimanual_pick_plate
bimanual_put_item_in_drawer bimanual_put_bottle_in_fridge
bimanual_handover_item bimanual_pick_laptop
bimanual_straighten_rope bimanual_sweep_to_dustpan
bimanual_lift_tray bimanual_handover_item_easy
bimanual_take_tray_out_of_oven

Project Structure

EnergyAction/
├── main.py                            # Entry point
├── setup.py                           # Package installation
├── modeling/
│   ├── ebm_compositionality/          # 🔑 Core EBM modules
│   │   ├── flow_to_energy.py          #   Flow Matching → Energy conversion
│   │   ├── energy_composer.py         #   E_left + E_right + E_coord composition
│   │   ├── ebm_bimanual_composer.py   #   Main EBM bimanual model
│   │   ├── bimanual_coordination_constraints.py  # Physics-informed constraints
│   │   └── panda_kinematics.py        #   Franka Panda robot kinematics
│   ├── encoder/                       # Vision-language encoders (CLIP backbone)
│   ├── noise_scheduler/               # DDPM / DDIM / Rectified Flow
│   ├── policy/                        # Base denoising policy networks
│   └── utils/                         # Transformer layers, attention, PE
├── datasets/                          # RLBench data loading (Zarr format)
├── utils/                             # Training loop, schedulers, EMA
├── pretrain_unimanual/                # Pretrained single-arm checkpoint
├── scripts/
│   └── rlbench/
│       ├── train_single_arm.sh        # Stage 1: Single-arm training
│       ├── train_ebm_bimanual_composer.sh  # Stage 2: EBM bimanual training
│       └── eval_ebm_bimanual_composer.sh   # Online evaluation
├── online_evaluation_rlbench/         # CoppeliaSim online evaluation
├── data_processing/                   # Raw data → Zarr conversion
├── data_generation/                   # RLBench data generation
└── instructions/                      # Task language instructions

Acknowledgement

Our implementation is mainly based on the following codebases. We gratefully thank the authors for their wonderful works.

Citation

@article{energyaction2026,
  title={EnergyAction: Energy-Based Compositional Flow Matching for Bimanual Manipulation},
  author={Song, Mingchen and Deng, Xiang and Wei, Jie and Jiang, Dongmei and Nie, Liqiang and Guan, Weili},
  journal={arXiv preprint arXiv:2603.20236},
  year={2026}
}

About

No description, website, or topics provided.

Resources

Stars

29 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages