MSG: Multi-Stream Generative Policies for Sample-Efficient Robotic Manipulation

Repository providing the source code for the paper "MSG: Multi-Stream Generative Policies for Sample-Efficient Robotic Manipulation", see the project website and video.

Please cite the paper as follows:

@ARTICLE{vonhartz2026msg,
  author={von Hartz, Jan Ole and Schweizer, Lukas and Boedecker, Joschka and Valada, Abhinav},
  journal={IEEE Robotics and Automation Letters}, 
  title={MSG: Multi-Stream Generative Policies for Sample-Efficient Robotic Manipulation}, 
  year={2026},
  volume={},
  number={},
  pages={1-8},
  doi={10.1109/LRA.2026.3682566}}

License

For academic usage, the code is released under the GPLv3 license. For any commercial purpose, please contact the authors.

Installation

Python

Developed with Python 3.10. Lower version will not work because of new syntax elements.

You can either use Conda or uv for installation.

Conda

• Create a Python env (using Conda) and install the latest Pytorch version.

• Install our RLBench fork form https://github.com/vonHartz/RLBench by following the instructions in that repo. After activating your RLBench conda env, run source rlbench_mode.sh to set up Coppelia and QT paths.

• Install the remaining requirements and this package via

pip install -r requirements.txt
pip install .

UV

• run uv sync to create the environment and install all dependencies.

• Edit rlbench_mode.sh by removing the first line (sourcing the conda env). Then source rlbench_mode.sh to set up Coppelia and QT paths.

MSG Quick Start

When using uv, prefix all commands with uv run.

Demonstration Data Collection (RLBench)

First we have to collect some demonstration data to train our Flow Matching models with.

python -m msg.collect_data_rlbench --config conf/collect_data/rlbench_expert.py --overwrite data_naming.task=StackWine

Training/Evaluation Examples

You can use the following commands to train/eval the ensemble-based MSG. For changing tasks and other parameters, please follow the instructions given in the respective config files. Please note that the suffix parameters during evaluation need to match the ones used during training!

Training (Single Object Task)

Start Frame:

python -m msg.behavior_cloning -f demos -t OpenDrawer --config conf/behavior_cloning/msg/rlbench/single_object_tasks/ee_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Target Frame:

python -m msg.behavior_cloning -f demos -t OpenDrawer --config conf/behavior_cloning/msg/rlbench/single_object_tasks/target_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Evaluation (Single Object Task)

python -m msg.evaluate -f demos -t OpenDrawer --config conf/evaluate/msg/rlbench/single_object_tasks/msg_ensemble.py

Training (Multi Object Task)

Subtask 1

Start Frame:

python -m msg.behavior_cloning -f demos -t StackWine --config conf/behavior_cloning/msg/rlbench/multi_object_tasks/subtraj1_ee_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Target Frame:

python -m msg.behavior_cloning -f demos -t StackWine --config conf/behavior_cloning/msg/rlbench/multi_object_tasks/subtraj1_object_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Subtask 2

Start Frame:

python -m msg.behavior_cloning -f demos -t StackWine --config conf/behavior_cloning/msg/rlbench/multi_object_tasks/subtraj2_ee_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Target Frame:

python -m msg.behavior_cloning -f demos -t StackWine --config conf/behavior_cloning/msg/rlbench/multi_object_tasks/subtraj2_target_frame.py  --overwrite training.wandb_mode=disabled training.seed=0

Evaluation (Multi Object Task)

python -m msg.evaluate -f demos -t StackWine --config conf/evaluate/msg/rlbench/multi_object_tasks/msg_ensemble.py```

Workflow

Entry points

For a full list of entry points, see setup.py.

Configs

Structure

Configs are stored as dataclass instances in individual modules in conf. The config structure roughly mirrors the code structure, though that's mostly for ease of use, as subconfigs are imported as modules anyway. Scripts usually need a config file and some command line args to supplement additional info. (Eg the task is a clarg so that one does not have to create one config file per task.) Furthermore, any config key can be overwritten from the command line using the --overwrite arg. Eg tapas-eval --config conf/evaluate/franka/banana.py -t Banana -f demos --overwrite wandb_mode=disabled policy.suffix=tx

Machine Config

Some config keys are machine specific, such as the data root. For this, I recommend placing a file called _machine.py in the conf folder and import it whenever needed. Here's an example of how the file may look like.

from omegaconf import MISSING

from msg.utils.misc import DataNamingConfig

data_naming_config = DataNamingConfig(
    feedback_type=MISSING, task=MISSING, data_root="data"
)


CUDA_PATH = "/usr/local/cuda-11.1"
LD_LIBRARY_PATH = ":".join(
    [CUDA_PATH + "/lib64", "/home/USERNAME/CoppeliaSim_Edu_V4_1_0_Ubuntu18_04"]
)

Then you can use from conf._machine import data_naming_config in any config file to use the correct machine specific data root without need to update your configs across machines.