Skip to content

async collect buffer for VLA RL #122

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yhnsu wants to merge 16 commits into
base: main
Choose a base branch
from
Open

async collect buffer for VLA RL #122

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
1a16fbb
async collect buffer for VLA RL
Feb 6, 2026
26246a5
vla policy wrapper
Feb 6, 2026
fd298f8
update
Feb 6, 2026
a4db089
Fix: clear buffer after get
yangchen73 Feb 27, 2026
c8c0919
Fix: add lock for get/is_full
yangchen73 Feb 27, 2026
71a5627
Fix: protect thread safety for on_step callback in Trainer
yangchen73 Feb 27, 2026
e73a9cc
Fix: removed repeated creation of normal distribution in ActorCritic
yangchen73 Feb 27, 2026
5dc18bf
Fix: include final smaller batch
yangchen73 Feb 27, 2026
da2eeeb
Remove buffer_type in train_config
yangchen73 Feb 27, 2026
b5dfd05
Sign unification
yangchen73 Feb 27, 2026
8f38083
Fix: correctly compute gae and set batch_size=[N, T] for vla_buffer
yangchen73 Mar 2, 2026
f75a43f
Merge remote-tracking branch 'origin/main' into yhn/rl_vla
yangchen73 Mar 3, 2026
130c05b
Merge remote-tracking branch 'origin/main' into yhn/rl_vla
yangchen73 Mar 9, 2026
c8140cf
Move policy adaptor from dexechain to embodichain
yangchen73 Mar 9, 2026
d8c30e1
reformate files
yangchen73 Mar 9, 2026
7ac335a
Step 32 instead of 1
yangchen73 Mar 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions configs/agents/rl/push_cube/train_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,15 @@
"gpu_id": 0,
"num_envs": 64,
"iterations": 1000,
"rollout_steps": 1024,
"buffer_size": 1024,
"enable_eval": true,
"num_eval_envs": 16,
"num_eval_episodes": 3,
"eval_freq": 2,
"eval_freq": 200,
"save_freq": 200,
Comment on lines +16 to 17
Copy link

Copilot AI Mar 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With the current Trainer logic, eval_freq/save_freq are interpreted in environment steps (global_step increments by buffer_size * num_envs each rollout). Values like 200 will almost never satisfy global_step % eval_freq == 0 (e.g., 64 envs * 1024 steps = 65536 per rollout). Either adjust the config to use step-based values that align with global_step, or update Trainer to a threshold-based schedule.

Copilot uses AI. Check for mistakes.
"use_wandb": false,
"wandb_project_name": "embodychain-push_cube",
"model_type": "standard",
"events": {
"eval": {
"record_camera": {
Expand All @@ -38,6 +39,7 @@
},
"policy": {
"name": "actor_critic",
"action_dim": 8,
"actor": {
"type": "mlp",
"network_cfg": {
Expand Down
61 changes: 61 additions & 0 deletions configs/agents/rl/stack_bowls/train_config_grpo.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
{
"trainer": {
"exp_name": "stack_bowls_vla_grpo",
"gym_config": "/root/workspace/research/embodichain/configs/gym/stack_bowls/gym_config.json",
"seed": 42,
"device": "cuda:0",
"headless": true,
"enable_rt": false,
"gpu_id": 0,
"num_envs": 8,
"iterations": 1000,
"buffer_size": 64,
"enable_eval": false,
"eval_freq": 0,
"save_freq": 100,
"use_wandb": false,
"wandb_project_name": "embodychain-stack_bowls",
"import_modules": [
"dexechain.lab.gym.envs.tasks.tableware.stack_bowls_v1.stack_bowls"
],
"model_type": "standard"
},
"policy": {
"name": "vla",
"action_dim": 14,
"vla_config": {
"model_path": "/root/workspace/output/stack_bowls/checkpoint-19000",
"instruction": "Stack the bowls.",
"inference_horizon": 32,
"action_std_init": 0.01,
"robot_type": "CobotMagic",
"gripper_open_value": 0.05,
"gripper_closed_value": 0.0,
"action_key_order": [
"left_armdelta_qpos",
"left_eefgripper",
"right_armdelta_qpos",
"right_eefgripper"
],
"model_config": {
"torch_dtype": "float32"
}
}
},
"algorithm": {
"name": "grpo",
"cfg": {
"learning_rate": 1e-5,
"n_epochs": 4,
"batch_size": 256,
"gamma": 0.99,
"clip_coef": 0.2,
"ent_coef": 0.001,
"kl_coef": 0.02,
"group_size": 4,
"eps": 1e-8,
"max_grad_norm": 1.0,
"truncate_at_first_done": true
Comment on lines +45 to +58
Copy link

Copilot AI Mar 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This config selects "algorithm": {"name": "grpo"}, but GRPO in embodichain/agents/rl/algo/grpo.py still uses the old collect_rollout() / update() (no-args) interface and the old policy.get_action() API. With the new Trainer/Policy TensorDict flow, this config will fail at runtime unless GRPO is migrated to update(rollout: TensorDict) (and collection moved to Collectors) or the algorithm name is changed to one that supports the new interface.

Copilot uses AI. Check for mistakes.
}
}
}
70 changes: 70 additions & 0 deletions configs/agents/rl/vla_example/train_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
{
"trainer": {
"exp_name": "vla_fine_tuning_ppo",
"gym_config": "configs/agents/rl/push_cube/gym_config.json",
"seed": 42,
"device": "cuda:0",
"headless": true,
"enable_rt": false,
"gpu_id": 0,
"num_envs": 32,
"iterations": 500,
"buffer_size": 2048,
"enable_eval": true,
"num_eval_envs": 8,
"num_eval_episodes": 3,
"eval_freq": 100,
"save_freq": 100,
"use_wandb": true,
"wandb_project_name": "embodychain-vla-training",
"model_type": "vla",
"events": {
"eval": {
"record_camera": {
"func": "record_camera_data_async",
"mode": "interval",
"interval_step": 1,
"params": {
"name": "main_cam",
"resolution": [640, 480],
"eye": [-1.4, 1.4, 2.0],
"target": [0, 0, 0],
"up": [0, 0, 1],
"intrinsics": [600, 600, 320, 240],
"save_path": "./outputs/videos/vla_eval"
}
}
}
}
},
"policy": {
"name": "vla",
"action_dim": 7,
"vla_config": {
"model_path": "checkpoints/pretrained_vla_model.pth",
"model_class": "vla_models.GPTVLAModel",
"model_config": {
"vision_encoder": "resnet50",
"language_model": "gpt2-medium",
"action_head_hidden_size": 512,
"freeze_vision_encoder": false,
"freeze_language_model": false
}
}
},
"algorithm": {
"name": "ppo",
"cfg": {
"learning_rate": 1e-5,
"n_epochs": 4,
"batch_size": 2048,
"gamma": 0.99,
"gae_lambda": 0.95,
"rollout_time_first": false,
"clip_coef": 0.2,
"ent_coef": 0.001,
"vf_coef": 0.5,
"max_grad_norm": 1.0
}
}
}
Loading
Loading