File size: 11,010 Bytes
567a9bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | diff --git a/pyproject.toml b/pyproject.toml
index 8bddca2..13a059b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@ dependencies = [
"imageio>=2.36.1",
"jax[cuda12]==0.5.3",
"jaxtyping==0.2.36",
- "lerobot[dataset]",
+ "lerobot==0.4.4",
"ml_collections==1.0.0",
"numpy>=1.22.4,<2.0.0",
"numpydantic>=1.6.6",
@@ -35,7 +35,7 @@ dependencies = [
"beartype==0.19.0",
"chex==0.1.90",
"treescope>=0.1.7",
- "transformers==5.5.4",
+ "transformers==4.53.2",
"rich>=14.0.0",
"polars>=1.30.0",
]
@@ -63,11 +63,11 @@ rlds = [
]
[tool.uv]
-override-dependencies = ["ml-dtypes==0.4.1", "tensorstore==0.1.74"]
+override-dependencies = ["ml-dtypes==0.4.1", "tensorstore==0.1.74", "numpy>=1.26,<3"]
[tool.uv.sources]
openpi-client = { workspace = true }
-lerobot = { git = "https://github.com/wensi-ai/lerobot", branch = "release/b1k" }
+
dlimp = { git = "https://github.com/kvablack/dlimp", rev = "ad72ce3a9b414db2185bc0b38461d4101a65477a" }
[tool.uv.workspace]
diff --git a/scripts/b1k/train_b1k.py b/scripts/b1k/train_b1k.py
index b53dc5c..0caac6d 100644
--- a/scripts/b1k/train_b1k.py
+++ b/scripts/b1k/train_b1k.py
@@ -467,6 +467,12 @@ def main(config: _config.TrainConfig):
reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos))
info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items())
pbar.write(f"Step {step}: {info_str}")
+ # pbar.write is a no-op here: openpi depends on tqdm_loggable, which
+ # swaps tqdm for a logger-backed bar that only forwards set_postfix.
+ # So the intended local fallback never reaches the job log, and with
+ # wandb disabled a run records no loss at all (jobs 2105-2110). Log it
+ # explicitly so metrics survive regardless of wandb.
+ logging.info("Step %d: %s", step, info_str)
wandb.log(reduced_info, step=step)
infos = []
if config.val_log_interval and step % config.val_log_interval == 0:
diff --git a/scripts/train_pytorch.py b/scripts/train_pytorch.py
index c7ddd2b..1f9aa97 100644
--- a/scripts/train_pytorch.py
+++ b/scripts/train_pytorch.py
@@ -429,6 +429,42 @@ def train_loop(config: _config.TrainConfig):
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,expandable_segments:True"
logging.info("Enabled memory optimizations for 8+ GPU training")
+ # Base weights and LoRA must both be in place *before* DDP wraps the model:
+ # DDP snapshots the parameter set at construction, and peft renames modules
+ # (base_model.model.*), so loading the plain safetensors state dict afterwards
+ # would not match. The original code loaded weights after the DDP block; that
+ # ordering only works without LoRA.
+ if config.pytorch_weight_path is not None:
+ _base_ckpt = os.path.join(config.pytorch_weight_path, "model.safetensors")
+ logging.info(f"Loading base weights from: {_base_ckpt}")
+ safetensors.torch.load_model(model, _base_ckpt)
+
+ if os.environ.get("OPENPI_LORA", "0") == "1":
+ from peft import LoraConfig, get_peft_model
+
+ # Follows RLinf's openpi recipe (rlinf/models/__init__.py): wrap only the
+ # PaliGemma VLM, leaving the action expert fully trainable. Their VLA runs
+ # use r=32 with lora_alpha == r and gaussian init.
+ rank = int(os.environ.get("OPENPI_LORA_RANK", "32"))
+ lora_config = LoraConfig(
+ r=rank,
+ lora_alpha=rank,
+ lora_dropout=0.0,
+ target_modules=[
+ "proj", "qkv", "fc1", "fc2", # vision
+ "q", "kv", "fc3", "out_proj", # projector
+ "q_proj", "k_proj", "v_proj", "o_proj", # llm attn
+ "gate_proj", "up_proj", "down_proj", "lm_head", # llm mlp
+ ],
+ init_lora_weights="gaussian",
+ )
+ _vlm = model.paligemma_with_expert.paligemma
+ model.paligemma_with_expert.paligemma = get_peft_model(_vlm, lora_config)
+ n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ n_all = sum(p.numel() for p in model.parameters())
+ logging.info(f"LoRA enabled (r={rank}): trainable {n_train/1e6:.1f}M / {n_all/1e9:.3f}B "
+ f"({100*n_train/n_all:.2f}%)")
+
if use_ddp:
model = torch.nn.parallel.DistributedDataParallel(
model,
@@ -438,15 +474,7 @@ def train_loop(config: _config.TrainConfig):
static_graph=world_size >= 8, # Enable for 8+ GPUs
)
- # Load weights from weight_loader if specified (for fine-tuning)
- if config.pytorch_weight_path is not None:
- logging.info(f"Loading weights from: {config.pytorch_weight_path}")
-
- model_path = os.path.join(config.pytorch_weight_path, "model.safetensors")
- safetensors.torch.load_model(
- (model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model), model_path
- )
- logging.info(f"Loaded PyTorch weights from {config.pytorch_weight_path}")
+ # (base weights are now loaded before the DDP/LoRA block above)
# Optimizer + learning rate schedule from config
warmup_steps = config.lr_schedule.warmup_steps
@@ -455,8 +483,11 @@ def train_loop(config: _config.TrainConfig):
end_lr = config.lr_schedule.decay_lr
# Create optimizer with config parameters
+ # Only optimise what LoRA left trainable; passing frozen params here would put
+ # them in the optimiser's param groups for no reason.
+ _trainable = [p for p in model.parameters() if p.requires_grad]
optim = torch.optim.AdamW(
- model.parameters(),
+ _trainable,
lr=peak_lr,
betas=(config.optimizer.b1, config.optimizer.b2),
eps=config.optimizer.eps,
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index ec6cd0f..9b3c8d9 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -772,6 +772,68 @@ _CONFIGS = [
assets_base_dir="./outputs/assets",
checkpoint_base_dir="./outputs/checkpoints",
),
+ # Local variant of pi05_b1k.
+ #
+ # `pi05_b1k` hardcodes the original author's dataset path, and DataConfig.base_config
+ # is `tyro.conf.Suppress`, so the `--data.base_config.dataset_root=...` override the
+ # challenge's baselines page documents is rejected by this code as an unrecognized
+ # option. A named config is the only way to point training at a different dataset
+ # without touching the upstream entry.
+ #
+ # Also switches the gemma variants to their LoRA forms, which makes
+ # Pi0Config.get_freeze_filter() freeze the base weights and train only the adapters.
+ TrainConfig(
+ name="pi05_b1k_lora_local",
+ model=pi0_config.Pi0Config(
+ action_horizon=32,
+ pi05=True,
+ paligemma_variant="gemma_2b_lora",
+ action_expert_variant="gemma_300m_lora",
+ ),
+ data=LeRobotB1KDataConfig(
+ repo_id="turning_on_radio",
+ base_config=DataConfig(
+ data_cls=_lerobot_compat.LeRobotDataset,
+ dataset_root="/data2/hyeongjinkim/behavior-challenge/data/b1k/turning_on_radio_v3.0",
+ prompt_from_task=True,
+ dataset_kwargs={"tolerance_s": 5e-4},
+ ),
+ robot_config_name="b1k/R1Pro",
+ ),
+ weight_loader=weight_loaders.CheckpointWeightLoader(
+ "/data2/hyeongjinkim/behavior-challenge/ckpt/pi05_base"
+ ),
+ save_interval=10_000,
+ num_train_steps=50_000,
+ assets_base_dir="./outputs/assets",
+ checkpoint_base_dir="./outputs/checkpoints",
+ ),
+ # PyTorch-path twin of the above. Same dataset, but the *base* gemma variants:
+ # LoRA is applied with peft in scripts/train_pytorch.py (OPENPI_LORA=1), which
+ # wraps standard nn.Linear modules and so must not see the JAX `*_lora` variants.
+ # Weights come from --pytorch-weight-path (safetensors), not weight_loader.
+ TrainConfig(
+ name="pi05_b1k_pytorch_local",
+ model=pi0_config.Pi0Config(action_horizon=32, pi05=True),
+ data=LeRobotB1KDataConfig(
+ repo_id="turning_on_radio",
+ base_config=DataConfig(
+ data_cls=_lerobot_compat.LeRobotDataset,
+ dataset_root="/data2/hyeongjinkim/behavior-challenge/data/b1k/turning_on_radio_v3.0",
+ prompt_from_task=True,
+ dataset_kwargs={"tolerance_s": 5e-4},
+ ),
+ robot_config_name="b1k/R1Pro",
+ ),
+ # Converted *inside the SFT venv*: that venv has transformers_replace applied,
+ # so the action expert carries AdaRMS (input_layernorm.dense.*). A checkpoint
+ # converted in the unpatched eval venv has plain RMSNorm and fails to load.
+ pytorch_weight_path="/data2/hyeongjinkim/behavior-challenge/ckpt/pi05_base_pytorch_sft",
+ save_interval=10_000,
+ num_train_steps=50_000,
+ assets_base_dir="./outputs/assets",
+ checkpoint_base_dir="./outputs/checkpoints",
+ ),
#
# Fine-tuning Libero configs.
#
diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py
index 2a8c13a..a7c4e40 100644
--- a/src/openpi/training/data_loader.py
+++ b/src/openpi/training/data_loader.py
@@ -165,13 +165,23 @@ def create_torch_dataset(
if repo_id == "fake":
return FakeDataset(model_config, num_samples=1024)
- dataset_meta = _lerobot_compat.LeRobotDatasetMetadata(repo_id)
+ # Honour dataset_root the way create_b1k_dataset does. Without it the PyTorch
+ # path ignores a local dataset and reaches for the HF hub, so a config that
+ # trains fine under train_b1k.py dies here with
+ # "FileNotFoundError: ~/.cache/huggingface/lerobot/<repo_id>/meta/info.json"
+ # and then a 404 on huggingface.co/api/datasets/<repo_id> (job 2128).
+ dataset_meta = _lerobot_compat.LeRobotDatasetMetadata(repo_id, root=data_config.dataset_root)
dataset = _lerobot_compat.LeRobotDataset(
data_config.repo_id,
+ root=data_config.dataset_root,
delta_timestamps={
key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys
},
- episodes=data_config.episodes_index,
+ # DataConfig has no `episodes_index` field — the original line referenced one
+ # and raised AttributeError for any config reaching this path (job 2129).
+ # getattr keeps the intent (optional episode subset) without inventing a field.
+ episodes=getattr(data_config, "episodes_index", None),
+ **data_config.dataset_kwargs,
)
if data_config.prompt_from_task:
|