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//meta/info.json" + # and then a 404 on huggingface.co/api/datasets/ (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: