Skip to content

README.md

Latest commit

 

History

History
288 lines (216 loc) · 7.09 KB

README.md

File metadata and controls

288 lines (216 loc) · 7.09 KB
Raw

Installation

Clone Repository

git clone https://github.com/FlagOpen/FlagScale.git
cd FlagScale/

Setup Conda Environment

Create a new conda environment for robotics training:

conda create -n flagscale-train python=3.12
conda activate flagscale-train

Install FlagScale and robotics dependencies:

cd FlagScale/
pip install . --verbose
pip install -r requirements/train/robotics/requirements.txt
pip install git+https://github.com/NVIDIA/Megatron-Energon.git@ab40226

Install Megatron:

mkdir -p /tmp
cd /tmp
git clone https://github.com/flagos-ai/Megatron-LM-FL.git
cd Megatron-LM-FL
pip install --no-build-isolation .[mlm,dev]

# add your path of FlagScale and the Megatron in FlagScale to PYTHONPATH
export PYTHONPATH=$PYTHONPATH:/xxx/FlagScale:/xxx/FlagScale/flagscale/train/

Training

Download Model

git lfs install

mkdir -p /models/BAAI/
cd /models/BAAI/
git clone https://huggingface.co/BAAI/RoboBrain-X0-Preview

If you don't have access to the international internet, download from modelscope.

mkdir -p /models/
cd /models/
modelscope download --model BAAI/RoboBrain-X0-Preview --local_dir BAAI/RoboBrain-X0-Preview

Prepare Dataset

FlagScale uses WebDataset format and Megatraon.Energon data loader, you need to process your data first.

There is a dataset processed: demo_0913_n2.

Download demo_0913_n2:

mkdir /tmp/datasets
cd /tmp/datasets
git clone https://gitee.com/hchnr/flag-scale.git
cd flag-scale
git checkout robotics_dataset

Move .jpg and .npy files from ./demo_0913_n2/deps to /:

mkdir -p /share/
cp -r ./demo_0913_n2/deps/* /

The directory structure of demo_0913_n2 is as follows:

  • build_dep.sh: Copy .npy and .jpg files from production environment to ./deps
  • demo_0913_n2.jsonl: Timesteps, including: task(str), images(.jpg), action(.npy), state(.npy)
  • deps: .npy and .jpg files
  • wds-2: Data in webdataset format (DP=2), generated by tools/datasets/vla/convert.py

If you need to make your own datasets, generate Data in webdataset format (DP=2) to ./demo_0913_n2/wds-2:

python tools/datasets/vla/convert.py \
    --dataset-root=./demo_0913_n2 \
    --output-root=./demo_0913_n2 \
    --json=demo_0913_n2.jsonl \
    --train-split 1 \
    --val-split 0 \
    --images-key=image \
    --videos-key=video \
    --vision-root='' \
    --shuffle-tars \
    --num-workers=1 \
    --max-samples-per-tar 100000 \
    --dp-size 2

Edit Config

cd FlagScale/
vim examples/robobrain_x0/conf/train/robobrain_x0.yaml

Change 2 fields:

  • data.data_path -> /tmp/datasets/flag-scale/demo_0913_n2/wds-1
  • data.tokenizer -> /models/BAAI/RoboBrain-X0-Preview

Start Training

cd FlagScale/
flagscale train robobrain_x0 --config ./examples/robobrain_x0/conf/train.yaml
# or
flagscale train robobrain_x0 -c ./examples/robobrain_x0/conf/train.yaml

Training with Lerobot Dataset

Note: For better performance, we recommend using WebDataset/Energon format (see Training section above). LeRobotDataset support is provided for convenience when working with existing LeRobot datasets.

LeRobotDataset Structure

A LeRobotDataset directory should have the following structure:

your_dataset/
├── data/
│   ├── chunk-000/
│   │   ├── file-000.parquet
│   │   ├── file-001.parquet
│   │   └── ...
│   └── ...
├── meta/
│   ├── info.json
│   ├── stats.json
│   ├── tasks.parquet
│   └── episodes/
│       └── ...
└── videos/  (optional, for video data)
    ├── observation.images.laptop/
    │   ├── chunk-000/
    │   │   ├── file-000.mp4
    │   │   └── ...
    │   └── ...
    └── ...

Download LeRobotDataset

You can download existing LeRobotDatasets from HuggingFace Hub:

# Example: Download aloha_mobile_cabinet dataset
pip install huggingface_hub
huggingface-cli download lerobot/aloha_mobile_cabinet --repo-type dataset --local-dir /datasets/lerobot/aloha_mobile_cabinet

Or use datasets from HuggingFace LeRobot collection.

Edit Config

FlagScale provides a pre-configured template for LeRobotDataset training:

cd FlagScale/
vim examples/robobrain_x0/conf/train/robobrain_x0_lerobot.yaml

Change the following fields according to your environment:

  • data.data_path: Path to your LeRobotDataset directory (e.g., /datasets/lerobot/aloha_mobile_cabinet)
  • data.tokenizer.tokenizer_path: Path to the RoboBrain-X0 model (e.g., /models/BAAI/RoboBrain-X0-Preview)
  • system.checkpoint.pretrained_checkpoint: Path to the pretrained checkpoint

Key configuration options:

data:
  # Path to your LeRobotDataset root directory
  data_path: /path/to/your/lerobot/dataset

  # Dataset type: use 'lerobot' for LeRobotDataset format
  dataset_type: lerobot

  # Video decoding backend: pyav (default), torchcodec, or video_reader
  video_backend: pyav

  # Number of bins for action discretization
  action_discretization_bins: 2048

Update train.yaml

Make sure train.yaml uses the lerobot configuration:

vim examples/robobrain_x0/conf/train.yaml

Set the default config to use lerobot:

defaults:
  - train: robobrain_x0_lerobot
  - _self_

Start Training

cd FlagScale/
flagscale train robobrain_x0 --config ./examples/robobrain_x0/conf/train.yaml
# or
flagscale train robobrain_x0 -c ./examples/robobrain_x0/conf/train.yaml

Supported Video Backends

For decoding video frames from LeRobotDataset, the following backends are supported:

  • pyav: Default backend, widely compatible
  • torchcodec: Faster decoding, requires torchcodec installation
  • video_reader: Torchvision's video reader backend

Serving

Download Tokenzier

mkdir -p /models/physical-intelligence/
cd /models/physical-intelligence/
git lfs install
git clone https://huggingface.co/physical-intelligence/fast

Edit Config

cd FlagScale/
vim examples/robobrain_x0/conf/serve/robobrain_x0.yaml

Change 3 fields:

  • engine_args.model_sub_task -> /models/BAAI/RoboBrain-X0-Preview
  • engine_args.port -> A port available in your env, for example: 5001
  • engine_args.tokenizer_path ->/models/physical-intelligence/fast

Run Serving

cd FlagScale/
flagscale serve robobrain_x0 --config ./examples/robobrain_x0/conf/serve.yaml
# or
flagscale serve robobrain_x0 -c ./examples/robobrain_x0/conf/serve.yaml

Test Server with Client

Download test images:

cd FlagScale/
wget https://gitee.com/hchnr/flag-scale/raw/robotics_dataset/orbbec_0_latest.jpg
wget https://gitee.com/hchnr/flag-scale/raw/robotics_dataset/orbbec_1_latest.jpg
wget https://gitee.com/hchnr/flag-scale/raw/robotics_dataset/orbbec_2_latest.jpg

Run client:

python examples/robobrain_x0/client_agilex.py  \
--host 127.0.0.1 \
--port 5001 \
--base-img orbbec_0_latest.jpg \
--left-wrist-img orbbec_1_latest.jpg \
--right-wrist-img orbbec_2_latest.jpg \
--num-steps 20