fix(profiling): wall_x OOM and xvla rename_map

- wall_x: switch to SGD optimizer + explicit scheduler overrides.
  The 4B-param model casts to bf16 internally, but AdamW's exp_avg/
  exp_avg_sq states blow past the 22 GB GPU. Same fix we applied to
  pi0/pi05/pi0_fast.
- xvla: fix rename_map. Dataset (libero_plus) exposes front/wrist
  image keys; the model expects image/image2. Previous map was
  direction-reversed and left the batch without any recognized
  image feature.

Made-with: Cursor
This commit is contained in:
Pepijn
2026-04-16 19:49:12 +02:00
parent d483dd4c4b
commit b81eef43c8
+12 -2
View File
@@ -152,7 +152,17 @@
"--policy.prediction_mode=diffusion",
"--policy.attn_implementation=eager",
"--policy.device=cuda",
"--batch_size=1"
"--batch_size=1",
"--use_policy_training_preset=false",
"--optimizer.type=sgd",
"--optimizer.lr=1e-5",
"--optimizer.weight_decay=0",
"--optimizer.grad_clip_norm=1.0",
"--scheduler.type=cosine_decay_with_warmup",
"--scheduler.peak_lr=1e-5",
"--scheduler.decay_lr=1e-6",
"--scheduler.num_warmup_steps=0",
"--scheduler.num_decay_steps=12"
]
},
"xvla": {
@@ -165,7 +175,7 @@
"--policy.empty_cameras=1",
"--policy.device=cuda",
"--batch_size=1",
"--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.image2\": \"observation.images.camera2\"}"
"--rename_map={\"observation.images.front\": \"observation.images.image\", \"observation.images.wrist\": \"observation.images.image2\"}"
]
}
}