Views
No views yet
pip install "lerobot[xvla]"select_action1import torch
2from lerobot.datasets.lerobot_dataset import LeRobotDataset
3from lerobot.policies.factory import make_pre_post_processors
4
5# Swap this import per-policy
6from lerobot.policies.xvla.modeling_xvla import XVLAPolicy
7
8# load a policy
9model_id = "lerobot/xvla-agibot-world" # <- swap checkpoint
10device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
12policy = XVLAPolicy.from_pretrained(model_id).to(device).eval()
13
14preprocess, postprocess = make_pre_post_processors(
15 policy.config,
16 model_id,
17 preprocessor_overrides={"device_processor": {"device": str(device)}},
18)
19# load a lerobotdataset (we will replace with a simpler dataset)
20dataset = LeRobotDataset("lerobot/libero")
21
22# pick an episode
23episode_index = 0
24
25# each episode corresponds to a contiguous range of frame indices
26from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
27to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
28
29# get a single frame from that episode (e.g. the first frame)
30frame_index = from_idx
31frame = dict(dataset[frame_index])
32
33batch = preprocess(frame)
34with torch.inference_mode():
35 pred_action = policy.select_action(batch)
36 # use your policy postprocess, this post process the action
37 # for instance unnormalize the actions, detokenize it etc..
38 pred_action = postprocess(pred_action)forward(...) to get a loss and then:1policy.train()
2batch = dict(dataset[0])
3batch = preprocess(batch)
4
5loss, outputs = policy.forward(batch)
6loss.backward()
7Notes:
- Some policies expose
policy(**batch)or return a dict; keep this snippet aligned with the policy API.- Use your trainer script (
lerobot-train) for full training loops.
1lerobot-train \
2 --dataset.repo_id=${HF_USER}/<dataset> \
3 --output_dir=./outputs/[RUN_NAME] \
4 --job_name=[RUN_NAME] \
5 --policy.repo_id=${HF_USER}/<desired_policy_repo_id> \
6 --policy.path=lerobot/[BASE_CHECKPOINT] \
7 --policy.dtype=bfloat16 \
8 --policy.device=cuda \
9 --steps=100000 \
10 --batch_size=4-policy.chunk_size=...-policy.n_action_steps=...-policy.max_action_tokens=...-policy.gradient_checkpointing=truerecord script from lerobot-record with a policy checkpoint as input, to run inference and evaluate your policy.lerobot-record \
--robot.type=so100_follower \
--robot.port=/dev/ttyACM1 \
--robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \
--robot.id=my_awesome_follower_arm \
--display_data=false \
--dataset.repo_id=${HF_USER}/eval_so100 \
--dataset.single_task="Put lego brick into the transparent box" \
# <- Teleop optional if you want to teleoperate in between episodes \
# --teleop.type=so100_leader \
# --teleop.port=/dev/ttyACM0 \
# --teleop.id=my_awesome_leader_arm \
--policy.path=${HF_USER}/my_policy