logo Mr.Robota 具身智能学习手册 中文
Chapter 10: Reproduction of RL Training for Heterogeneous Robots with UniLab + MotrixSim

Baseline Evaluation of Low-Level Manipulations in Home Organization Tasks

Paper | Official Website | Model | Dataset | Supplementary Materials


Installation and Setup

  1. Installation Environment

    bashconda create -n mshab python=3.9
    conda activate mshab
    
  2. Install ManiSkill3 and project code

    bashgit clone https://github.com/haosulab/ManiSkill.git -b mshab --single-branch
    pip install -e ManiSkill
    pip install -e .  # 可选安装训练/开发依赖:pip install -e .[train,dev]
    
  3. Download the required dataset

    bashfor dataset in ycb ReplicaCAD ReplicaCADRearrange; do python -m mani_skill.utils.download_asset "$dataset"; done
    
  4. Import Environment

    pythonimport mshab.envs
    

[Optional] Checkpoints, dataset, and data generation

  • Models and dataset can be obtained from HuggingFace (approximately 490GB):

    bashhuggingface-cli login
    
    huggingface-cli download arth-shukla/mshab_checkpoints --local-dir mshab_checkpoints
    
    export MS_ASSET_DIR="~/.maniskill"
    export MSHAB_DATASET_DIR="$MS_ASSET_DIR/data/scene_datasets/replica_cad_dataset/rearrange-dataset"
    
    huggingface-cli download --repo-type dataset arth-shukla/MS-HAB-TidyHouse --local-dir "$MSHAB_DATASET_DIR/tidy_house"
    huggingface-cli download --repo-type dataset arth-shukla/MS-HAB-PrepareGroceries --local-dir "$MSHAB_DATASET_DIR/prepare_groceries"
    huggingface-cli download --repo-type dataset arth-shukla/MS-HAB-SetTable --local-dir "$MSHAB_DATASET_DIR/set_table"
    
  • Generate data using a script (customizable trajectory filtering criteria):

    bashbash scripts/gen_dataset.sh
    

Quick Start

MS-HAB provides:

  • Evaluation environment SequentialTask-v0: Used for complete task evaluation (supports sub-task chains)

  • Training environment [Name]SubtaskTrain-v0: Used to train individual skill policies, such as PickSubtaskTrain-v0

Example code is as follows:

pythonimport gymnasium as gym
from mani_skill.vector.wrappers.gymnasium import ManiSkillVectorEnv
import mshab.envs
from mshab.envs.planner import plan_data_from_file

task = "tidy_house"
subtask = "pick"
split = "train"

plan_data = plan_data_from_file(REARRANGE_DIR / "task_plans" / task / subtask / split / "all.json")
spawn_data_fp = REARRANGE_DIR / "spawn_data" / task / subtask / split / "spawn_data.pt"

env = gym.make(
    f"{subtask.capitalize()}SubtaskTrain-v0",
    num_envs=252,
    obs_mode="rgbd",
    sim_backend="gpu",
    robot_uids="fetch",
    control_mode="pd_joint_delta_pos",
    reward_mode="normalized_dense",
    render_mode="rgb_array",
    max_episode_steps=200,
    task_plans=plan_data.plans,
    scene_builder_cls=plan_data.dataset,
    spawn_data_fp=spawn_data_fp,
)

venv = ManiSkillVectorEnv(env, max_episode_steps=1000, ignore_terminations=True)
obs, info = venv.reset()

Scene and Task Customization

  • Scenario: ReplicaCAD provides 84 apartment scenarios (63 training, 21 validation)

  • Task Plan: Defines the order and objectives of tasks, such as “Pick → Place”

  • Generate data: The initial state of robots and objects, pre-generated by the script (for stability during GPU simulation)

Simple Customization

  • Only some task plans can be loaded for debugging.

  • The success/failure criteria can be adjusted via task_cfgs (e.g., the range of the target placement)

Advanced Customization

  • The task chain can be extended (it must be logical, such as avoiding “Pick → Pick”).

  • You can change objects/scenes, etc. by modifying the configuration and using SceneBuilder.


Training

Support multiple algorithms (SAC, PPO, behavior cloning, diffusion policy):

bashbash scripts/train_[algo].sh

Example:

bashpython -m mshab.train_sac \
  mshab_checkpoints/rl/tidy_house/pick/all/config.yml \
  algo.gamma=0.8

Restore Training

bashpython -m mshab.train_ppo \
  mshab_checkpoints/rl/tidy_house/pick/all/config.yml \
  model_ckpt=mshab_checkpoints/rl/tidy_house/pick/all/policy.pt \
  algo.lr=1e-3

Evaluation

Data must be downloaded or generated first. The following script can be used for the evaluation of full tasks or sub-tasks:

bashbash scripts/evaluate_sequential_task.sh