VLA · Phase 7 · VLA Implementation and Experimentation
7.4 — Fine-Tuning on Robot Demonstrations
Adapting a pretrained model to a new robot or task with your own data.
A pretrained policy is a starting point. This lesson adapts one to a new robot with 24 demonstrations in three ways and compares what each trains and how well it works.
Claims about papers, models and textbooks carry numbered references, listed with what each one is cited for in the References. Code marked “Toy model” is our own illustration, not the cited system.
Learning objectives
After this lesson, you should be able to:
- Explain why pretrained policies are fine-tuned.
- Compare full, head-only and LoRA fine-tuning by trained parameters.
- Read the test loss before and after adaptation.
- List practical rules for a fine-tuning run.
- Start from a pretrained checkpoint
- Prepare demonstrations from the new robot
- Choose what to train: all, head or low-rank corrections
- Train with a small learning rate, watching validation
- Evaluate on the robot or simulator
The one-sentence idea
Start from what is known and train only what has to change.
Fewer trained parameters, less memory, less forgetting.
Why fine-tune
A released VLA was trained on someone else's robots, cameras and tasks. To use it on yours, you continue training on your own demonstrations: fine-tuning. OpenVLA's authors report that it outperforms the from-scratch method Diffusion Policy by 20.4% in fine-tuning experiments on multi-task settings with multiple objects and language grounding,[1] and that it can be fine-tuned on consumer GPUs via modern low-rank adaptation methods.[1] Octo likewise reports fine-tuning to new sensory inputs and action spaces within a few hours on consumer GPUs.[2]
This lesson uses PyTorch (these cells run on the course's run server) and a tiny network, so you can watch the whole process: pretrain, then adapt in three ways, and compare.
A pretrained network and a new robot
The “pretrained” model is a small network trained on plenty of data for one robot. The “new robot” has a slightly different mapping from observation to action and only 24 demonstrations. We start by pretraining and measuring how the pretrained model does on the new robot before any adaptation:
import torch, torch.nn as nn
torch.manual_seed(0)
def robot_data(n, shift):
x = torch.randn(n, 8)
w = torch.linspace(-1, 1, 8)
y = torch.tanh(x @ w) + shift * x[:, 0] # robot B differs from robot A by the shift term
return x, y.unsqueeze(1)
net = nn.Sequential(nn.Linear(8, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
def train(model, params, x, y, steps, lr):
opt = torch.optim.Adam(params, lr=lr)
for _ in range(steps):
opt.zero_grad(); loss = ((model(x) - y) ** 2).mean(); loss.backward(); opt.step()
def mse(model, x, y):
with torch.no_grad(): return float(((model(x) - y) ** 2).mean())
xa, ya = robot_data(2000, shift=0.0) # plenty of data from robot A
train(net, net.parameters(), xa, ya, steps=400, lr=3e-3)
xb_test, yb_test = robot_data(2000, shift=0.6)
print("pretrained, robot A test loss:", round(mse(net, *robot_data(500, 0.0)), 3))
print("pretrained, robot B test loss:", round(mse(net, xb_test, yb_test), 3))Output:
pretrained, robot A test loss: 0.0 pretrained, robot B test loss: 0.355
Toy model. A small network and synthetic data. It shows the mechanics of adaptation, not the behaviour of a real VLA.
The pretrained model is good on its own robot and clearly worse on the new one. Now adapt with 24 demonstrations.
Three ways to adapt
(1) Full fine-tuning: update every weight. (2) Head only: freeze everything except the last layer. (3) Low-rank adaptation (LoRA): freeze the weights and learn a small low-rank correction to a layer (lesson 6.1). Compare the number of trained parameters and the test loss on robot B:
import copy, torch, torch.nn as nn
torch.manual_seed(0)
def robot_data(n, shift):
x = torch.randn(n, 8); w = torch.linspace(-1, 1, 8)
return x, (torch.tanh(x @ w) + shift * x[:, 0]).unsqueeze(1)
def mse(model, x, y):
with torch.no_grad(): return float(((model(x) - y) ** 2).mean())
def train(model, params, x, y, steps, lr):
opt = torch.optim.Adam(params, lr=lr)
for _ in range(steps):
opt.zero_grad(); loss = ((model(x) - y) ** 2).mean(); loss.backward(); opt.step()
# the same pretraining as before
net = nn.Sequential(nn.Linear(8, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
xa, ya = robot_data(2000, 0.0); train(net, net.parameters(), xa, ya, 400, 3e-3)
class LoRALinear(nn.Module):
def __init__(self, base, rank):
super().__init__()
self.base = base
for p in self.base.parameters(): p.requires_grad = False
self.A = nn.Parameter(torch.randn(base.in_features, rank) * 0.1)
self.B = nn.Parameter(torch.zeros(rank, base.out_features)) # starts at zero: the model begins unchanged
def forward(self, x): return self.base(x) + x @ self.A @ self.B
torch.manual_seed(1)
x_small, y_small = robot_data(24, 0.6) # 24 demonstrations on robot B
x_test, y_test = robot_data(2000, 0.6)
def count(params): return sum(p.numel() for p in params)
full = copy.deepcopy(net)
train(full, full.parameters(), x_small, y_small, 150, 1e-3)
head = copy.deepcopy(net)
for p in head.parameters(): p.requires_grad = False
for p in head[4].parameters(): p.requires_grad = True
head_params = [p for p in head.parameters() if p.requires_grad]
train(head, head_params, x_small, y_small, 150, 1e-2)
lora = copy.deepcopy(net)
lora[0] = LoRALinear(lora[0], rank=2)
for p in lora[2].parameters(): p.requires_grad = False
for p in lora[4].parameters(): p.requires_grad = False
lora_params = [p for p in lora.parameters() if p.requires_grad]
train(lora, lora_params, x_small, y_small, 150, 1e-2)
print("before adapting: test loss", round(mse(net, x_test, y_test), 3))
print("full fine-tuning: test loss", round(mse(full, x_test, y_test), 3), "| trained parameters:", count(full.parameters()))
print("head only: test loss", round(mse(head, x_test, y_test), 3), "| trained parameters:", count(head_params))
print("LoRA (rank 2, layer 1): test loss", round(mse(lora, x_test, y_test), 3), "| trained parameters:", count(lora_params))Output:
before adapting: test loss 0.36 full fine-tuning: test loss 0.037 | trained parameters: 4801 head only: test loss 0.022 | trained parameters: 65 LoRA (rank 2, layer 1): test loss 0.072 | trained parameters: 144
All three beat the unadapted model (0.36), using 24 demonstrations. In this toy the head-only version came out best (0.022), which is not a general result: here the new robot differs from the old one only at the output, so changing only the output suffices. When the difference is spread through the network, more of it has to change. The trade-offs are the usual ones: full fine-tuning can change the most but touches every weight (and, with big models, needs the most memory and risks forgetting, lesson 5.4); the cheaper methods train far fewer parameters.
Rules of thumb
Hold out whole episodes for validation (lesson 3.2). Watch the validation loss and stop when it turns up: with 24 demonstrations, a network can memorise them. Keep the learning rate small for pretrained weights. Evaluate on the robot or simulator, not only on loss (lesson 3.4). And keep the pretrained checkpoint, the data, the split and the statistics together so the run can be repeated (lesson 7.2).
Common misconceptions
Misconception 1 “Fine-tuning needs as much data as pretraining.”
Correction A few dozen demonstrations moved the toy model a long way.
Misconception 2 “Low loss means a working robot.”
Correction Evaluate in closed loop (lesson 3.4).
Misconception 3 “LoRA is always as good as full fine-tuning.”
Correction It depends; the paper reports parity in its setting.
Engineering takeaways
- Fine-tune on your robot’s demonstrations.
- Three scopes: all, head, low-rank.
- Hold out whole episodes.
- Evaluate on the robot.
Change as little as you can, as much as you must.
Knowledge check
Try each question first, then reveal the explanation.
Question 1
Why initialise LoRA’s B matrix to zero?
Explanation
So the adapted model starts identical to the pretrained one.
Question 2
Which of the three methods trains the fewest parameters?
Explanation
LoRA (in the experiment), followed by head only.
Question 3
Why watch the validation loss with 24 demonstrations?
Explanation
The network can memorise so few examples.
References
Numbered notes in the text, like [1], point here. Each entry says what it is cited for and how that was checked. Papers and books can be revised or misread, so follow the links before relying on a detail. Code marked “Toy model” is our own illustration of an idea, not the code of the cited system.
(2024). OpenVLA: An open-source vision-language-action model. Conference on Robot Learning (CoRL) 2024. arXiv:2406.09246
Cited for:
- In fine-tuning experiments it outperforms the from-scratch imitation learning method Diffusion Policy by 20.4% on multi-task settings with multiple objects and language grounding. Checked against the paper’s abstract.
- OpenVLA can be fine-tuned on consumer GPUs via modern low-rank adaptation methods and served efficiently via quantization without a hit to downstream success rate. Checked against the paper’s abstract.
(2024). Octo: An open-source generalist robot policy. RSS 2024 / arXiv. arXiv:2405.12213
Cited for:
- Can be effectively finetuned to robot setups with new sensory inputs and action spaces within a few hours on standard consumer GPUs; tested across 9 robotic platforms. Checked against the paper’s abstract.
What’s next
Next: training a small policy from scratch to see every step.