Views
No views yet
Base Planning Model: The model was trained on general datasets in Stages 1–2 and on the Robotic Planning dataset in Stage 3, which is designed for Planning prediction.A-LoRA for Affordance: Based on the Base Planning Model, Stage 4 involves LoRA-based training with our Affordance dataset to predict affordance.T-LoRA for Trajectory: Based on the Base Planning Model, Stage 4 involves LoRA-based training with our Trajectory dataset to predict trajectory.
| Models | Checkpoint | Description |
|---|---|---|
| Planning Model | 🤗 Planning CKPTs | Used for Planning prediction in our paper |
| Affordance (A-LoRA) | 🤗 Affordance CKPTs | Used for Affordance prediction in our paper |
| Trajectory (T-LoRA) | 🤗 Trajectory CKPTs | Used for Trajectory prediction in our paper |
1# clone repo.
2git clone https://github.com/FlagOpen/RoboBrain.git
3cd RoboBrain
4# build conda env.
5conda create -n robobrain python=3.10
6conda activate robobrain
7pip install -r requirements.txt1# Modify datasets for Stage 4_traj, please refer to:
2- yaml_path: scripts/train/yaml/stage_4_trajectory.yaml1{
2 "id": 0,
3 "image": [
4 "shareRobot/trajectory/images/rtx_frames_success_0/10_utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds#episode_2/frame_0.png"
5 ],
6 "conversations": [
7 {
8 "from": "human",
9 "value": "<image>\nYou are a robot using the joint control. The task is \"reach for the cloth\". Please predict up to 10 key trajectory points to complete the task. Your answer should be formatted as a list of tuples, i.e. [[x1, y1], [x2, y2], ...], where each tuple contains the x and y coordinates of a point."
10 },
11 {
12 "from": "gpt",
13 "value": "[[0.781, 0.305], [0.688, 0.344], [0.570, 0.344], [0.492, 0.312]]"
14 }
15 ]
16},1# Training on Stage 4_traj:
2bash scripts/train/stage_4_0_resume_finetune_lora_t.sh1# Planning Model
2python model/llava_utils/convert_robobrain_to_hf.py --model_dir /path/to/original/checkpoint/ --dump_path /path/to/output/
3# A-LoRA & T-RoRA
4python model/llava_utils/convert_lora_weights_to_hf.py --model_dir /path/to/original/checkpoint/ --dump_path /path/to/output/1# please refer to https://github.com/FlagOpen/RoboBrain
2from inference import SimpleInference
3model_id = "BAAI/RoboBrain"
4lora_id = "BAAI/RoboBrain-LoRA-Affordance"
5model = SimpleInference(model_id, lora_id)
6# Example 1:
7prompt = "You are a robot using the joint control. The task is \"reach for the cloth\". Please predict up to 10 key trajectory points to complete the task. Your answer should be formatted as a list of tuples, i.e. [[x1, y1], [x2, y2], ...], where each tuple contains the x and y coordinates of a point."
8image = "./assets/demo/trajectory_1.jpg"
9pred = model.inference(prompt, image, do_sample=False)
10print(f"Prediction: {pred}")
11'''
12 Prediction: [[0.781, 0.305], [0.688, 0.344], [0.570, 0.344], [0.492, 0.312]]
13'''
14# Example 2:
15prompt = "You are a robot using the joint control. The task is \"reach for the grapes\". Please predict up to 10 key trajectory points to complete the task. Your answer should be formatted as a list of tuples, i.e. [[x1, y1], [x2, y2], ...], where each tuple contains the x and y coordinates of a point."
16image = "./assets/demo/trajectory_2.jpg"
17pred = model.inference(prompt, image, do_sample=False)
18print(f"Prediction: {pred}")
19'''
20 Prediction: [[0.898, 0.352], [0.766, 0.344], [0.625, 0.273], [0.500, 0.195]]
21'''1@article{ji2025robobrain,
2 title={RoboBrain: A Unified Brain Model for Robotic Manipulation from Abstract to Concrete},
3 author={Ji, Yuheng and Tan, Huajie and Shi, Jiayu and Hao, Xiaoshuai and Zhang, Yuan and Zhang, Hengyuan and Wang, Pengwei and Zhao, Mengdi and Mu, Yao and An, Pengju and others},
4 journal={arXiv preprint arXiv:2502.21257},
5 year={2025}
6}