ElasticVLA：面向实时场景的VLA模型推理与调度系统

本项目是北京大学 OSLab Embodied AI System 项目的一部分。

项目简介

ElasticVLA 旨在解决大型视觉-语言-动作（VLA）模型在端侧设备上推理延迟高、难以满足实时控制需求的问题。本项目在开源 VLA 大模型 OpenVLA （moojink/openvla-oft） 的基础上实现了 ElasticVLA，引入了多出口架构（Multi-Exit）、语义适配器（Semantic Adapter）、模型路由器（Router）以及实时调度器（Scheduler），实现了在保证控制精度的前提下，动态调整模型推理深度，并利用节省的计算资源执行背景任务（如传感器数据处理等）。

---

环境配置

# Create and activate conda environment
conda create -n openvla-oft python=3.10 -y
conda activate openvla-oft

# Install PyTorch
# Use a command specific to your machine: https://pytorch.org/get-started/locally/
pip3 install torch torchvision torchaudio

# Clone openvla-oft repo and pip install to download dependencies
git clone https://git.pku.edu.cn/2301111924/ElasticVLA
cd ElasticVLA
pip install -e .

# Install Flash Attention 2 for training (https://github.com/Dao-AILab/flash-attention)
#   =>> If you run into difficulty, try `pip cache remove flash_attn` first
pip install packaging ninja
ninja --version; echo $?  # Verify Ninja --> should return exit code "0"
pip install "flash-attn==2.5.5" --no-build-isolation

模型训练

LIBERO 训练环境配置

首先按照以上步骤配置环境。

然后你需要从以下链接克隆并配置 LIBERO 运行库：

git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git
pip install -e LIBERO
pip install -r experiments/robot/libero/libero_requirements.txt  # From openvla-oft base dir

对于高版本 pytorch ，你需要修改 LIBERO 库中的 libero/libero/benchmark/__init__.py 文件，将 init_states = torch.load(init_states_path) 修改为 init_states = torch.load(init_states_path,weights_only=False)。

接下来，你需要下载 LIBERO 数据集以进行模型训练：

git clone git@hf.co:datasets/openvla/modified_libero_rlds

基础模型下载

本项目使用 OpenVLA-OFT 作为基础模型进行微调。你可以从以下链接下载预训练好的 OpenVLA-OFT 模型：

https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-spatial-object-goal-10

1. 多出口模型微调

使用以下命令启动多出口模型的微调训练：

 torchrun --standalone --nnodes 1 --nproc-per-node 1 vla-scripts/finetune.py \
  --vla_path BASIC_VLA_MODEL_PATH \
  --data_root_dir DATA_ROOT_PATH \
  --run_root_dir CHECKPOINTS_AND_LOG_SAVE_DIR \
  --dataset_name "libero_spatial_no_noops" \
  --use_l1_regression True \
  --train_action_head_only True \
  --exit_block_indices "7,15,23,31" \
  --exit_loss_weights "1,1,1,1" \
  --use_film False \
  --num_images_in_input 2 \
  --use_proprio True \
  --batch_size 8 \
  --grad_accumulation_steps 1 \
  --max_steps MAX_STEPS  \
  --num_steps_before_decay NUM_STEPS_BEFORE_DECAY \
  --save_freq SAVE_FREQ \
  --image_aug False \
  --use_lora True \
  --lora_rank 32 \
  --wandb_entity YOUR_WANDB_ENTITY \
  --wandb_project YOUR_WANDB_PROJECT

2. 语义适配器（Semantic Adapter）微调

torchrun --standalone --nnodes 1 --nproc-per-node 1  vla-scripts/finetune.py \
  --vla_path BASIC_VLA_MODEL_PATH \
  --data_root_dir DATA_ROOT_PATH \
  --run_root_dir CHECKPOINTS_AND_LOG_SAVE_DIR \
  --dataset_name "libero_spatial_no_noops" \
  --train_semantic_adapter_only True \
  --use_l1_regression True \
  --exit_block_indices "7,15,23,31" \
  --exit_loss_weights "1,1,1,1" \
  --use_film False \
  --num_images_in_input 2 \
  --use_proprio True \
  --adapter_source_layer 15 \
  --adapter_target_layer 31 \
  --batch_size 8 \
  --grad_accumulation_steps 2 \
  --max_steps MAX_STEPS  \
  --num_steps_before_decay NUM_STEPS_BEFORE_DECAY \
  --save_freq SAVE_FREQ \
  --wandb_log_freq 1 \
  --image_aug False \
  --use_val_set False \
  --wandb_entity YOUR_WANDB_ENTITY \
  --wandb_project YOUR_WANDB_PROJECT

3. 模型路由器（Router） 微调

torchrun --standalone --nnodes 1 --nproc-per-node 1  vla-scripts/finetune.py \
  --vla_path BASIC_VLA_MODEL_PATH \
  --data_root_dir DATA_ROOT_PATH \
  --run_root_dir CHECKPOINTS_AND_LOG_SAVE_DIR \
  --dataset_name "libero_spatial_no_noops" \
  --train_router_only True \
  --router_label_threshold 0.1 \
  --use_l1_regression True \
  --exit_block_indices "7,15,23,31" \
  --exit_loss_weights "1,1,1,1" \
  --use_film False \
  --num_images_in_input 2 \
  --use_proprio True \
  --adapter_source_layer 15 \
  --adapter_target_layer 31 \
  --learning_rate 1e-5 \
  --batch_size 16 \
  --grad_accumulation_steps 4 \
  --max_steps MAX_STEPS  \
  --num_steps_before_decay NUM_STEPS_BEFORE_DECAY \
  --save_freq SAVE_FREQ \
  --wandb_log_freq 1 \
  --image_aug False \
  --use_val_set False \
  --wandb_entity YOUR_WANDB_ENTITY \
  --wandb_project YOUR_WANDB_PROJECT

---

评估与实时调度实验

在完成上述训练后，你可以使用 run_libero_eval.py 脚本在 LIBERO 任务集上评估模型，并测试 ElasticVLA 的实时调度能力。

运行实验

我们提供了一个自动化测试脚本示例 scripts/run.sh。

chmod +x scripts/run.sh
./scripts/run.sh

核心调度参数说明

在 run_libero_eval.py 中，我们新增了以下与实时调度相关的参数：

• --use_adapter_inference True: 启用 Semantic Adapter 推理。
• --use_scheduler True: 启用实时调度器。
• --scheduler_baseline_mode False: 是否使用 Baseline 模式。如果为 True，调度器将悲观地预留最慢出口的时间，不利用提前退出的空闲时间。
• --scheduler_mode "utility_knapsack": 调度器模式。utility_knapsack 表示使用 0/1 背包算法在 VLA 推理和背景任务之间进行联合优化。
• --scheduler_loss_penalty 20.0: 背包算法中，预测误差的惩罚系数（$\lambda \backslash lambda$）。
• --router_threshold 0.1: Router 判定 FAST 或 SLOW 的阈值。
• --control_cycle_time_ms 250: 系统的控制周期时间预算（毫秒）。
• --exit_layer_indices_str "7,15,23,31": 模型所有可用的出口层索引。
• --exit_costs_ms "15:120,31:220": 各个出口层的预估耗时（毫秒）。格式为 层号:耗时。
• --fast_exit_layer 15: Router 判定为 FAST 时使用的出口层。
• --slow_exit_layer 31: Router 判定为 SLOW 时使用的出口层。
• --use_background_tasks True: 是否启用背景任务流。
• --background_tasks "cam1:30:5:0:1000; lidar1:50:10:200:800": 模拟的背景任务流。格式为 任务ID:耗时(ms):价值:到达时间(ms):过期时间(ms)，多个任务用分号 ; 隔开。
• --measure_performance True: 启用性能指标测量。
• --save_performance_metrics True: 将性能指标和调度器日志保存为 JSON 文件。
• --performance_log_dir "./performance_logs": 日志保存目录。

实验报告

详见 evaluation_report.md