| |
| |
| |
| |
| @@ -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] |
| |
| |
| |
| |
| @@ -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: |
| |
| |
| |
| |
| @@ -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, |
| |
| |
| |
| |
| @@ -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. |
| # |
| |
| |
| |
| |
| @@ -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: |
|
|
