| """ |
| Physics-Grounded Pretraining for Diffusion Policy using THE WELL. |
| |
| Streams diverse physics simulation data from HuggingFace to pretrain |
| the temporal prediction backbone before fine-tuning on robot demonstrations. |
| |
| Architecture: |
| THE WELL physics data → Spatiotemporal Encoder → Temporal Prediction Head |
| ↓ (transfer) |
| Robot demonstrations → Same Encoder → Diffusion Policy → Actions |
| """ |
|
|
| import math |
| import logging |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Optional |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader, IterableDataset |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| |
| |
| |
|
|
| @dataclass |
| class PhysicsPretrainConfig: |
| |
| datasets: list[str] = field(default_factory=lambda: [ |
| "polymathic-ai/gray_scott", |
| "polymathic-ai/acoustic_scattering", |
| "polymathic-ai/rayleigh_benard", |
| ]) |
| streaming: bool = True |
|
|
| |
| hidden_dim: int = 256 |
| n_layers: int = 4 |
| n_heads: int = 8 |
| context_frames: int = 4 |
| prediction_horizon: int = 16 |
| spatial_patch_size: int = 8 |
|
|
| |
| batch_size: int = 32 |
| learning_rate: float = 3e-4 |
| weight_decay: float = 1e-5 |
| num_steps: int = 200_000 |
| warmup_steps: int = 5_000 |
| gradient_clip: float = 1.0 |
| mixed_precision: bool = True |
| seed: int = 42 |
|
|
| |
| checkpoint_dir: str = "./checkpoints/physics_pretrain" |
| checkpoint_every: int = 10_000 |
| eval_every: int = 5_000 |
|
|
| |
| hf_repo: str = "arc-ai/diffusion-policy-physics" |
| push_to_hub: bool = True |
| push_every: int = 50_000 |
|
|
|
|
| @dataclass |
| class FinetuneConfig: |
| pretrained_path: str = "" |
| dataset_repo: str = "arc-ai/sim-demonstrations" |
| obs_dim: int = 20 |
| action_dim: int = 7 |
| action_horizon: int = 16 |
|
|
| batch_size: int = 256 |
| learning_rate: float = 1e-4 |
| num_steps: int = 500_000 |
| warmup_steps: int = 2_000 |
| freeze_encoder_steps: int = 50_000 |
| gradient_clip: float = 1.0 |
| mixed_precision: bool = True |
|
|
| |
| domain_randomization: bool = True |
| noise_sensor_std: float = 0.05 |
| noise_actuator_std: float = 0.02 |
| mass_range: tuple[float, float] = (0.7, 1.5) |
| friction_range: tuple[float, float] = (0.5, 2.0) |
| gravity_range: tuple[float, float] = (9.0, 10.5) |
|
|
|
|
| |
| |
| |
|
|
| class WellStreamingDataset(IterableDataset): |
| """Streams physics trajectories from THE WELL via HuggingFace.""" |
|
|
| def __init__(self, config: PhysicsPretrainConfig): |
| self.config = config |
| self.datasets = config.datasets |
| self.context_frames = config.context_frames |
| self.prediction_horizon = config.prediction_horizon |
| self.patch_size = config.spatial_patch_size |
|
|
| def _load_stream(self): |
| from datasets import load_dataset, interleave_datasets |
|
|
| streams = [] |
| for ds_name in self.datasets: |
| try: |
| ds = load_dataset(ds_name, split="train", streaming=True) |
| streams.append(ds) |
| logger.info(f"Streaming from {ds_name}") |
| except Exception as e: |
| logger.warning(f"Failed to load {ds_name}: {e}") |
|
|
| if not streams: |
| raise RuntimeError("No datasets loaded") |
|
|
| return interleave_datasets(streams) |
|
|
| def _extract_trajectory(self, sample): |
| """Extract spatiotemporal trajectory from WELL sample. |
| |
| THE WELL format: B × T × H [× W [× D]] × C |
| We extract temporal windows for prediction. |
| """ |
| if "trajectory" in sample: |
| data = np.array(sample["trajectory"], dtype=np.float32) |
| elif "fields" in sample: |
| data = np.array(sample["fields"], dtype=np.float32) |
| else: |
| keys = [k for k in sample.keys() if k not in ("metadata", "id", "label")] |
| if keys: |
| data = np.array(sample[keys[0]], dtype=np.float32) |
| else: |
| return None |
|
|
| if data.ndim < 2: |
| return None |
|
|
| T = data.shape[0] |
| total_needed = self.context_frames + self.prediction_horizon |
| if T < total_needed: |
| return None |
|
|
| return data |
|
|
| def _patchify(self, spatial_data: np.ndarray) -> np.ndarray: |
| """Convert spatial field to patch embeddings.""" |
| if spatial_data.ndim == 1: |
| return spatial_data |
|
|
| if spatial_data.ndim == 2: |
| H, C = spatial_data.shape |
| n_patches = H // self.patch_size |
| if n_patches == 0: |
| return spatial_data.reshape(-1) |
| patches = spatial_data[:n_patches * self.patch_size].reshape(n_patches, self.patch_size, C) |
| return patches.mean(axis=1).reshape(-1) |
|
|
| if spatial_data.ndim == 3: |
| H, W, C = spatial_data.shape |
| pH = H // self.patch_size |
| pW = W // self.patch_size |
| if pH == 0 or pW == 0: |
| return spatial_data.reshape(-1)[:256] |
| patches = spatial_data[:pH*self.patch_size, :pW*self.patch_size].reshape( |
| pH, self.patch_size, pW, self.patch_size, C |
| ) |
| return patches.mean(axis=(1, 3)).reshape(-1)[:256] |
|
|
| return spatial_data.reshape(-1)[:256] |
|
|
| def __iter__(self): |
| stream = self._load_stream() |
|
|
| for sample in stream: |
| trajectory = self._extract_trajectory(sample) |
| if trajectory is None: |
| continue |
|
|
| T = trajectory.shape[0] |
| total_needed = self.context_frames + self.prediction_horizon |
|
|
| start = np.random.randint(0, T - total_needed + 1) |
| window = trajectory[start:start + total_needed] |
|
|
| context_patches = [] |
| for t in range(self.context_frames): |
| patch = self._patchify(window[t]) |
| context_patches.append(patch) |
|
|
| target_patches = [] |
| for t in range(self.context_frames, total_needed): |
| patch = self._patchify(window[t]) |
| target_patches.append(patch) |
|
|
| context = np.stack(context_patches, axis=0) |
| target = np.stack(target_patches, axis=0) |
|
|
| yield { |
| "context": torch.from_numpy(context), |
| "target": torch.from_numpy(target), |
| } |
|
|
|
|
| |
| |
| |
|
|
| class SinusoidalEmbedding(nn.Module): |
| def __init__(self, dim: int): |
| super().__init__() |
| self.dim = dim |
|
|
| def forward(self, t: torch.Tensor) -> torch.Tensor: |
| half = self.dim // 2 |
| emb = math.log(10000) / (half - 1) |
| emb = torch.exp(torch.arange(half, device=t.device) * -emb) |
| emb = t.unsqueeze(-1) * emb.unsqueeze(0) |
| return torch.cat([emb.sin(), emb.cos()], dim=-1) |
|
|
|
|
| class PhysicsTemporalEncoder(nn.Module): |
| """Encoder that learns spatiotemporal dynamics from physics simulations. |
| |
| Pretrained on THE WELL, then transferred to robot policy. |
| """ |
|
|
| def __init__(self, config: PhysicsPretrainConfig): |
| super().__init__() |
| self.config = config |
| dim = config.hidden_dim |
|
|
| self.input_proj = nn.Sequential( |
| nn.LazyLinear(dim), |
| nn.SiLU(), |
| nn.Linear(dim, dim), |
| ) |
|
|
| self.time_embed = SinusoidalEmbedding(dim) |
| self.time_proj = nn.Sequential( |
| nn.Linear(dim, dim), |
| nn.SiLU(), |
| nn.Linear(dim, dim), |
| ) |
|
|
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=dim, |
| nhead=config.n_heads, |
| dim_feedforward=dim * 4, |
| dropout=0.1, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.transformer = nn.TransformerEncoder( |
| encoder_layer, num_layers=config.n_layers |
| ) |
|
|
| self.output_norm = nn.LayerNorm(dim) |
|
|
| def forward(self, context: torch.Tensor, timesteps: Optional[torch.Tensor] = None): |
| """ |
| Args: |
| context: (B, T, D) - temporal context frames (patchified) |
| timesteps: (B,) - optional diffusion timesteps |
| |
| Returns: |
| (B, T, hidden_dim) - encoded temporal representation |
| """ |
| B, T, _ = context.shape |
|
|
| x = self.input_proj(context) |
|
|
| if timesteps is not None: |
| t_emb = self.time_proj(self.time_embed(timesteps.float())) |
| x = x + t_emb.unsqueeze(1) |
|
|
| pos = torch.arange(T, device=x.device).float() |
| pos_emb = self.time_embed(pos).unsqueeze(0).expand(B, -1, -1) |
| x = x + pos_emb[:, :T, :x.shape[-1]] |
|
|
| x = self.transformer(x) |
| x = self.output_norm(x) |
| return x |
|
|
|
|
| |
| |
| |
|
|
| class PhysicsPredictionHead(nn.Module): |
| """Predicts future physics states from encoded context.""" |
|
|
| def __init__(self, config: PhysicsPretrainConfig): |
| super().__init__() |
| dim = config.hidden_dim |
| self.horizon = config.prediction_horizon |
|
|
| self.predictor = nn.Sequential( |
| nn.Linear(dim, dim * 2), |
| nn.SiLU(), |
| nn.Linear(dim * 2, dim * 2), |
| nn.SiLU(), |
| ) |
| self.output_proj = nn.LazyLinear(1) |
|
|
| def forward(self, encoded: torch.Tensor, target_dim: int): |
| """ |
| Args: |
| encoded: (B, T_context, dim) |
| target_dim: output dimension per timestep |
| Returns: |
| (B, horizon, target_dim) |
| """ |
| pooled = encoded.mean(dim=1) |
| x = self.predictor(pooled) |
|
|
| if self.output_proj.in_features == 1: |
| self.output_proj = nn.Linear( |
| x.shape[-1], self.horizon * target_dim |
| ).to(x.device) |
|
|
| out = self.output_proj(x) |
| return out.reshape(-1, self.horizon, target_dim) |
|
|
|
|
| |
| |
| |
|
|
| class PhysicsDiffusionPolicy(nn.Module): |
| """Diffusion Policy using physics-pretrained temporal encoder.""" |
|
|
| def __init__(self, obs_dim: int, action_dim: int, config: PhysicsPretrainConfig): |
| super().__init__() |
| self.obs_dim = obs_dim |
| self.action_dim = action_dim |
| self.action_horizon = config.prediction_horizon |
| dim = config.hidden_dim |
|
|
| self.encoder = PhysicsTemporalEncoder(config) |
|
|
| self.obs_proj = nn.Sequential( |
| nn.Linear(obs_dim, dim), |
| nn.SiLU(), |
| nn.Linear(dim, dim), |
| ) |
|
|
| self.noise_pred = nn.Sequential( |
| nn.Linear(dim * 2, dim * 2), |
| nn.SiLU(), |
| nn.Linear(dim * 2, dim * 2), |
| nn.SiLU(), |
| nn.Linear(dim * 2, action_dim * self.action_horizon), |
| ) |
|
|
| n_steps = 100 |
| betas = torch.linspace(1e-4, 0.02, n_steps) |
| alphas = 1.0 - betas |
| alphas_cumprod = torch.cumprod(alphas, dim=0) |
|
|
| self.register_buffer("betas", betas) |
| self.register_buffer("alphas_cumprod", alphas_cumprod) |
| self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod)) |
| self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1 - alphas_cumprod)) |
| self.n_steps = n_steps |
|
|
| def forward(self, obs: torch.Tensor, actions: torch.Tensor, timesteps: torch.Tensor): |
| """Training forward: predict noise added to actions. |
| |
| Args: |
| obs: (B, T_obs, obs_dim) - observation history |
| actions: (B, horizon, action_dim) - ground truth action chunk |
| timesteps: (B,) - diffusion timesteps |
| """ |
| encoded = self.encoder(self.obs_proj(obs), timesteps) |
| context = encoded.mean(dim=1) |
|
|
| noise = torch.randn_like(actions) |
| B = actions.shape[0] |
| t = timesteps.long() |
|
|
| sqrt_alpha = self.sqrt_alphas_cumprod[t].view(B, 1, 1) |
| sqrt_one_minus = self.sqrt_one_minus_alphas_cumprod[t].view(B, 1, 1) |
| noisy_actions = sqrt_alpha * actions + sqrt_one_minus * noise |
|
|
| noisy_flat = noisy_actions.reshape(B, -1) |
| combined = torch.cat([context, noisy_flat[:, :context.shape[-1]]], dim=-1) |
|
|
| noise_pred = self.noise_pred(combined) |
| noise_pred = noise_pred.reshape(B, self.action_horizon, self.action_dim) |
|
|
| return F.mse_loss(noise_pred, noise) |
|
|
| @torch.no_grad() |
| def predict(self, obs: torch.Tensor, n_inference_steps: int = 10) -> torch.Tensor: |
| """DDIM inference: generate action chunk from observation. |
| |
| Args: |
| obs: (B, T_obs, obs_dim) |
| n_inference_steps: number of denoising steps |
| Returns: |
| (B, horizon, action_dim) |
| """ |
| B = obs.shape[0] |
| device = obs.device |
|
|
| actions = torch.randn(B, self.action_horizon, self.action_dim, device=device) |
|
|
| step_indices = torch.linspace(self.n_steps - 1, 0, n_inference_steps).long().to(device) |
|
|
| for step_t in step_indices: |
| t_batch = step_t.expand(B) |
| encoded = self.encoder(self.obs_proj(obs), t_batch.float()) |
| context = encoded.mean(dim=1) |
|
|
| noisy_flat = actions.reshape(B, -1) |
| combined = torch.cat([context, noisy_flat[:, :context.shape[-1]]], dim=-1) |
|
|
| noise_pred = self.noise_pred(combined) |
| noise_pred = noise_pred.reshape(B, self.action_horizon, self.action_dim) |
|
|
| alpha = self.alphas_cumprod[step_t] |
| alpha_prev = self.alphas_cumprod[max(step_t - self.n_steps // n_inference_steps, 0)] |
|
|
| pred_x0 = (actions - torch.sqrt(1 - alpha) * noise_pred) / torch.sqrt(alpha) |
| pred_x0 = pred_x0.clamp(-1, 1) |
|
|
| actions = torch.sqrt(alpha_prev) * pred_x0 + torch.sqrt(1 - alpha_prev) * noise_pred |
|
|
| return actions |
|
|
|
|
| |
| |
| |
|
|
| class PhysicsPretrainer: |
| """Pretrains temporal encoder on THE WELL physics data.""" |
|
|
| def __init__(self, config: PhysicsPretrainConfig): |
| self.config = config |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| self.encoder = PhysicsTemporalEncoder(config).to(self.device) |
| self.pred_head = PhysicsPredictionHead(config).to(self.device) |
|
|
| params = list(self.encoder.parameters()) + list(self.pred_head.parameters()) |
| self.optimizer = torch.optim.AdamW( |
| params, lr=config.learning_rate, weight_decay=config.weight_decay |
| ) |
|
|
| self.scaler = torch.amp.GradScaler("cuda", enabled=config.mixed_precision) |
| self.step = 0 |
|
|
| def _get_lr(self) -> float: |
| if self.step < self.config.warmup_steps: |
| return self.config.learning_rate * self.step / self.config.warmup_steps |
| progress = (self.step - self.config.warmup_steps) / ( |
| self.config.num_steps - self.config.warmup_steps |
| ) |
| return self.config.learning_rate * 0.5 * (1 + math.cos(math.pi * progress)) |
|
|
| def _update_lr(self): |
| lr = self._get_lr() |
| for pg in self.optimizer.param_groups: |
| pg["lr"] = lr |
|
|
| def train(self): |
| """Run pretraining loop streaming from THE WELL.""" |
| logger.info(f"Starting physics pretraining on {self.device}") |
| logger.info(f"Datasets: {self.config.datasets}") |
| logger.info(f"Steps: {self.config.num_steps}") |
|
|
| dataset = WellStreamingDataset(self.config) |
| dataloader = DataLoader( |
| dataset, batch_size=self.config.batch_size, num_workers=2, pin_memory=True |
| ) |
|
|
| checkpoint_dir = Path(self.config.checkpoint_dir) |
| checkpoint_dir.mkdir(parents=True, exist_ok=True) |
|
|
| running_loss = 0.0 |
| log_interval = 100 |
|
|
| for batch in dataloader: |
| if self.step >= self.config.num_steps: |
| break |
|
|
| context = batch["context"].to(self.device) |
| target = batch["target"].to(self.device) |
|
|
| self._update_lr() |
|
|
| with torch.amp.autocast("cuda", enabled=self.config.mixed_precision): |
| encoded = self.encoder(context) |
| target_flat = target.reshape(target.shape[0], target.shape[1], -1) |
| target_dim = target_flat.shape[-1] |
| predicted = self.pred_head(encoded, target_dim) |
|
|
| min_horizon = min(predicted.shape[1], target_flat.shape[1]) |
| min_dim = min(predicted.shape[2], target_flat.shape[2]) |
| loss = F.mse_loss( |
| predicted[:, :min_horizon, :min_dim], |
| target_flat[:, :min_horizon, :min_dim], |
| ) |
|
|
| self.optimizer.zero_grad() |
| self.scaler.scale(loss).backward() |
| self.scaler.unscale_(self.optimizer) |
| nn.utils.clip_grad_norm_( |
| list(self.encoder.parameters()) + list(self.pred_head.parameters()), |
| self.config.gradient_clip, |
| ) |
| self.scaler.step(self.optimizer) |
| self.scaler.update() |
|
|
| running_loss += loss.item() |
| self.step += 1 |
|
|
| if self.step % log_interval == 0: |
| avg_loss = running_loss / log_interval |
| lr = self._get_lr() |
| logger.info( |
| f"Step {self.step}/{self.config.num_steps} | " |
| f"Loss: {avg_loss:.6f} | LR: {lr:.2e}" |
| ) |
| running_loss = 0.0 |
|
|
| if self.step % self.config.checkpoint_every == 0: |
| self._save_checkpoint(checkpoint_dir / f"step_{self.step}.pt") |
|
|
| if self.step % self.config.push_every == 0 and self.config.push_to_hub: |
| self._push_to_hub() |
|
|
| self._save_checkpoint(checkpoint_dir / "final.pt") |
| logger.info(f"Pretraining complete at step {self.step}") |
| return self.encoder |
|
|
| def _save_checkpoint(self, path: Path): |
| torch.save({ |
| "step": self.step, |
| "encoder_state": self.encoder.state_dict(), |
| "pred_head_state": self.pred_head.state_dict(), |
| "optimizer_state": self.optimizer.state_dict(), |
| "config": self.config, |
| }, path) |
| logger.info(f"Saved checkpoint: {path}") |
|
|
| def _push_to_hub(self): |
| try: |
| from huggingface_hub import HfApi |
| api = HfApi() |
| path = Path(self.config.checkpoint_dir) / f"step_{self.step}.pt" |
| api.upload_file( |
| path_or_fileobj=str(path), |
| path_in_repo=f"checkpoints/step_{self.step}.pt", |
| repo_id=self.config.hf_repo, |
| repo_type="model", |
| ) |
| logger.info(f"Pushed step {self.step} to {self.config.hf_repo}") |
| except Exception as e: |
| logger.warning(f"Hub push failed: {e}") |
|
|
|
|
| |
| |
| |
|
|
| class DomainRandomizer: |
| """Applies domain randomization during fine-tuning.""" |
|
|
| def __init__(self, config: FinetuneConfig): |
| self.config = config |
|
|
| def __call__(self, obs: torch.Tensor, actions: torch.Tensor): |
| if not self.config.domain_randomization: |
| return obs, actions |
|
|
| B = obs.shape[0] |
|
|
| |
| obs = obs + torch.randn_like(obs) * self.config.noise_sensor_std |
|
|
| |
| actions = actions + torch.randn_like(actions) * self.config.noise_actuator_std |
|
|
| |
| mass_scale = torch.empty(B, 1, 1, device=obs.device).uniform_( |
| *self.config.mass_range |
| ) |
| if obs.shape[-1] >= 14: |
| obs[:, :, 7:14] = obs[:, :, 7:14] / mass_scale |
|
|
| return obs, actions |
|
|
|
|
| class PolicyFinetuner: |
| """Fine-tunes physics-pretrained encoder for robot manipulation.""" |
|
|
| def __init__(self, pretrained_encoder: PhysicsTemporalEncoder, config: FinetuneConfig): |
| self.config = config |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| pretrain_config = PhysicsPretrainConfig( |
| hidden_dim=pretrained_encoder.config.hidden_dim, |
| n_layers=pretrained_encoder.config.n_layers, |
| n_heads=pretrained_encoder.config.n_heads, |
| prediction_horizon=config.action_horizon, |
| ) |
|
|
| self.policy = PhysicsDiffusionPolicy( |
| obs_dim=config.obs_dim, |
| action_dim=config.action_dim, |
| config=pretrain_config, |
| ).to(self.device) |
|
|
| |
| self.policy.encoder.load_state_dict(pretrained_encoder.state_dict(), strict=False) |
| logger.info("Loaded pretrained encoder weights") |
|
|
| self.randomizer = DomainRandomizer(config) |
| self.optimizer = torch.optim.AdamW( |
| self.policy.parameters(), lr=config.learning_rate |
| ) |
| self.scaler = torch.amp.GradScaler("cuda", enabled=config.mixed_precision) |
| self.step = 0 |
|
|
| def _freeze_encoder(self, freeze: bool): |
| for p in self.policy.encoder.parameters(): |
| p.requires_grad = not freeze |
|
|
| def train(self, dataloader: DataLoader): |
| """Fine-tune on robot demonstrations with domain randomization.""" |
| logger.info("Starting fine-tuning with domain randomization") |
| self._freeze_encoder(True) |
|
|
| running_loss = 0.0 |
| for batch in dataloader: |
| if self.step >= self.config.num_steps: |
| break |
|
|
| if self.step == self.config.freeze_encoder_steps: |
| self._freeze_encoder(False) |
| logger.info(f"Unfreezing encoder at step {self.step}") |
|
|
| obs = batch["obs"].to(self.device) |
| actions = batch["actions"].to(self.device) |
|
|
| obs, actions = self.randomizer(obs, actions) |
|
|
| t = torch.randint(0, self.policy.n_steps, (obs.shape[0],), device=self.device) |
|
|
| with torch.amp.autocast("cuda", enabled=self.config.mixed_precision): |
| loss = self.policy(obs, actions, t.float()) |
|
|
| self.optimizer.zero_grad() |
| self.scaler.scale(loss).backward() |
| self.scaler.unscale_(self.optimizer) |
| nn.utils.clip_grad_norm_(self.policy.parameters(), self.config.gradient_clip) |
| self.scaler.step(self.optimizer) |
| self.scaler.update() |
|
|
| running_loss += loss.item() |
| self.step += 1 |
|
|
| if self.step % 100 == 0: |
| logger.info(f"Step {self.step} | Loss: {running_loss/100:.6f}") |
| running_loss = 0.0 |
|
|
| return self.policy |
|
|
|
|
| |
| |
| |
|
|
| def setup_hf_repo(repo_id: str, repo_type: str = "model"): |
| """Create HuggingFace repo if it doesn't exist.""" |
| from huggingface_hub import HfApi, create_repo |
|
|
| api = HfApi() |
| try: |
| create_repo(repo_id, repo_type=repo_type, exist_ok=True) |
| logger.info(f"HF repo ready: {repo_id}") |
| except Exception as e: |
| logger.warning(f"Repo setup: {e}") |
| return api |
|
|
|
|
| def upload_model(model: nn.Module, repo_id: str, filename: str = "model.pt"): |
| """Upload model checkpoint to HuggingFace.""" |
| import tempfile |
| from huggingface_hub import HfApi |
|
|
| api = HfApi() |
| with tempfile.NamedTemporaryFile(suffix=".pt") as f: |
| torch.save(model.state_dict(), f.name) |
| api.upload_file( |
| path_or_fileobj=f.name, |
| path_in_repo=filename, |
| repo_id=repo_id, |
| repo_type="model", |
| ) |
| logger.info(f"Uploaded {filename} to {repo_id}") |
|
|
|
|
| def upload_dataset(data_dir: str, repo_id: str): |
| """Upload demonstration dataset to HuggingFace.""" |
| from huggingface_hub import HfApi |
|
|
| api = HfApi() |
| api.upload_folder( |
| folder_path=data_dir, |
| repo_id=repo_id, |
| repo_type="dataset", |
| ) |
| logger.info(f"Uploaded dataset to {repo_id}") |
|
|
|
|
| |
| |
| |
|
|
| def run_physics_pretraining(config_path: Optional[str] = None): |
| """Full pipeline: pretrain on physics → fine-tune on robot demos.""" |
| import yaml |
|
|
| if config_path: |
| with open(config_path) as f: |
| cfg = yaml.safe_load(f) |
| pretrain_cfg = PhysicsPretrainConfig(**{ |
| k: v for k, v in cfg.get("pretraining", {}).items() |
| if k in PhysicsPretrainConfig.__dataclass_fields__ |
| }) |
| else: |
| pretrain_cfg = PhysicsPretrainConfig() |
|
|
| |
| logger.info("=" * 60) |
| logger.info("PHASE 1: Physics Pretraining (THE WELL)") |
| logger.info("=" * 60) |
|
|
| pretrainer = PhysicsPretrainer(pretrain_cfg) |
| encoder = pretrainer.train() |
|
|
| |
| logger.info("=" * 60) |
| logger.info("PHASE 2: Fine-tuning on Robot Demonstrations") |
| logger.info("=" * 60) |
|
|
| finetune_cfg = FinetuneConfig( |
| pretrained_path=str(Path(pretrain_cfg.checkpoint_dir) / "final.pt"), |
| ) |
|
|
| finetuner = PolicyFinetuner(encoder, finetune_cfg) |
|
|
| |
| |
| logger.info("Fine-tuning requires demonstration dataloader — provide via API") |
|
|
| return finetuner.policy |
|
|
|
|
| if __name__ == "__main__": |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") |
| run_physics_pretraining() |
|
|
