""" 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__) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- @dataclass class PhysicsPretrainConfig: # THE WELL datasets (streamed from HuggingFace) datasets: list[str] = field(default_factory=lambda: [ "polymathic-ai/gray_scott", "polymathic-ai/acoustic_scattering", "polymathic-ai/rayleigh_benard", ]) streaming: bool = True # Architecture 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 # Training 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 # Checkpointing checkpoint_dir: str = "./checkpoints/physics_pretrain" checkpoint_every: int = 10_000 eval_every: int = 5_000 # HuggingFace 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 during fine-tuning 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) # --------------------------------------------------------------------------- # THE WELL Streaming Dataset # --------------------------------------------------------------------------- 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), } # --------------------------------------------------------------------------- # Physics-Aware Temporal Encoder # --------------------------------------------------------------------------- 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 # --------------------------------------------------------------------------- # Prediction Head (for pretraining) # --------------------------------------------------------------------------- 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) # will adapt to target dim 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) # --------------------------------------------------------------------------- # Diffusion Policy with Physics-Pretrained Encoder # --------------------------------------------------------------------------- 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 # --------------------------------------------------------------------------- # Pretraining Loop # --------------------------------------------------------------------------- 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}") # --------------------------------------------------------------------------- # Fine-tuning on Robot Demonstrations # --------------------------------------------------------------------------- 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] # Sensor noise obs = obs + torch.randn_like(obs) * self.config.noise_sensor_std # Actuator noise actions = actions + torch.randn_like(actions) * self.config.noise_actuator_std # Mass scaling (affects dynamics — simulated via obs perturbation) 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) # Transfer pretrained weights 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 # --------------------------------------------------------------------------- # HuggingFace Hub Utilities # --------------------------------------------------------------------------- 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}") # --------------------------------------------------------------------------- # Main Entry Point # --------------------------------------------------------------------------- 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() # Phase 1: Pretrain on THE WELL logger.info("=" * 60) logger.info("PHASE 1: Physics Pretraining (THE WELL)") logger.info("=" * 60) pretrainer = PhysicsPretrainer(pretrain_cfg) encoder = pretrainer.train() # Phase 2: Fine-tune on robot demonstrations 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) # Load robot demonstration data (from HF or local) # In production, this streams from arc-ai/sim-demonstrations 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()