# MuJoCo Scene Designs and Code Architecture --- ## Scene 1: The Indistinguishable Pair (Chapter 1) Two visually identical cubes with different physics. This is the single most important demo because it makes the argument in 5 seconds. ### MuJoCo XML Sketch ```xml ``` ### What You Measure Push both cubes with identical force. Record: - Displacement over time - Final position - Contact forces The demo shows: identical visual input → completely different physical outcomes → video prediction fails, field + material estimation succeeds. --- ## Scene 2: The Tabletop (Chapters 2, 3, 4) The workhorse scene. 5 objects with varied materials on a table. Used for SDF reconstruction, material estimation, object tracking, and mixture-of-experts. ### Objects | Object | Shape | Material Properties | Purpose | |---|---|---|---| | Wooden block | Box | E=10GPa, ρ=600, μ=0.4, e=0.3 | Baseline rigid object | | Rubber ball | Sphere | E=0.01GPa, ρ=1100, μ=0.8, e=0.7 | High friction + bounce | | Metal cylinder | Cylinder | E=200GPa, ρ=7800, μ=0.2, e=0.5 | Heavy, low friction | | Plastic cup | Cylinder (hollow) | E=3GPa, ρ=1200, μ=0.35, e=0.4 | Container (for future pouring) | | Sponge | Box | E=0.001GPa, ρ=30, μ=0.6, e=0.1 | Deformable (soft body in MuJoCo) | ### Camera Setup 4 cameras to support viewpoint experiments: ```xml ``` --- ## Scene 3: The Occlusion Test (Chapter 3) Two identical objects. An occluder wall in the middle. One object is pushed behind the wall and comes back. ``` Time 0: [A] |wall| [B] ← A and B visible, distinct Time 1: [A→] |wall| [B] ← A moving toward wall Time 2: |wall| [B] ← A fully occluded Time 3: [←A] |wall| [B] ← A re-emerges. Is it still A? ``` For the field-only model: it has no persistent identity, so after occlusion it might assign A's latent to B or vice versa. The graph model maintains the node for A with existence probability decaying slightly but not zeroing out, and recovers identity on re-emergence via appearance + position matching. --- ## Scene 4: The Stack (Chapters 3, 6) Three blocks stacked. Tests support relation inference and prediction when a block is removed. ``` [C] ← top [B] ← middle [A] ← bottom (on table) ``` Remove B → C should fall. The graph knows (A supports B) and (B supports C). Removing B breaks the chain. The field predicts where C lands based on geometry + dynamics. --- ## Code Architecture ``` hybrid_world_model/ │ ├── envs/ # MuJoCo environments │ ├── base_env.py # Common: step, render, get_obs, get_ground_truth │ ├── indistinguishable.py # Scene 1 │ ├── tabletop.py # Scene 2 │ ├── occlusion.py # Scene 3 │ └── stack.py # Scene 4 │ ├── model/ │ ├── field.py # Neural field MLP │ │ class ObjectField(nn.Module): │ │ """Maps (x_local, z_i) → (sdf_mu, sdf_sigma, material_params)""" │ │ def __init__(self, latent_dim=128, hidden=256, layers=4): │ │ def forward(self, x_local, z_i): │ │ → returns dict: { │ │ 'sdf_mu': ..., 'sdf_sigma': ..., │ │ 'friction_mu': ..., 'friction_sigma': ..., │ │ 'density_mu': ..., 'density_sigma': ..., │ │ 'stiffness_mu': ..., 'stiffness_sigma': ... │ │ } │ │ │ ├── mixture.py # Mixture of experts │ │ class MixtureField(nn.Module): │ │ """Global field = Σ w_i * ObjectField_i + w_0 * BackgroundField""" │ │ def __init__(self, object_field, gating_network, n_objects): │ │ def query(self, x_world, object_poses, object_latents): │ │ # 1. Transform x_world to each object frame │ │ # 2. Query each local field │ │ # 3. Compute gating weights │ │ # 4. Return mixture distribution │ │ → returns dict with mixture predictions │ │ │ ├── gating.py # Gating network │ │ class GatingNetwork(nn.Module): │ │ """SDF-based prior + learned residual → ownership weights""" │ │ def forward(self, sdf_values, x_world, object_latents): │ │ → returns weights [w_0, w_1, ..., w_N], sum to 1 │ │ │ ├── graph.py # Object graph │ │ class ObjectGraph: │ │ """Nodes: (pose, latent, existence). Edges: (contact_prob, type)""" │ │ def __init__(self, max_objects=10): │ │ def update(self, observations): # Update poses, latents, contacts │ │ def get_contacts(self): # Return active contact pairs │ │ def get_support_chain(self, obj_i): # Trace support relations │ │ │ ├── encoder.py # Observation → state │ │ class WorldModelEncoder(nn.Module): │ │ """RGB + depth → updated (field, graph) state""" │ │ def __init__(self, vision_backbone='resnet18'): │ │ def forward(self, rgb, depth, prev_state): │ │ # 1. Extract visual features │ │ # 2. Segment objects (or use GT masks in sim) │ │ # 3. Update per-object latents z_i │ │ # 4. Update poses T_i │ │ # 5. Update contact probabilities │ │ → returns WorldState(field, graph) │ │ │ └── dynamics.py # State prediction │ class DynamicsModel(nn.Module): │ """Predict S_{t+1} from S_t and action a_t""" │ def __init__(self): │ def forward(self, state, action): │ # 1. GNN message passing on graph │ # 2. Update object latents and poses │ # 3. Update contact probabilities │ → returns predicted WorldState │ ├── training/ │ ├── data_collector.py # Run MuJoCo, collect trajectories │ ├── train_field.py # Train SDF field on GT data │ ├── train_encoder.py # Train encoder on observations │ ├── train_dynamics.py # Train dynamics on trajectory data │ └── train_full.py # End-to-end training │ ├── evaluation/ │ ├── sdf_metrics.py # SDF MAE, IoU, Chamfer, calibration │ ├── material_metrics.py # Material param estimation error │ ├── tracking_metrics.py # Identity tracking accuracy │ └── planning_metrics.py # Grasp/push success rate │ ├── planning/ │ ├── cem.py # Cross-entropy method planner │ ├── grasp_planner.py # Grasp force selection from field queries │ └── push_planner.py # Push planning using dynamics model │ ├── visualization/ │ ├── field_vis.py # SDF slices, isosurfaces, uncertainty heatmaps │ ├── graph_vis.py # Graph topology visualization │ ├── ownership_vis.py # Gating weight heatmaps │ └── material_convergence.py # Distribution evolution plots │ ├── demos/ # Gradio demo apps (one per chapter) │ ├── demo_ch1_why_not_video.py │ ├── demo_ch2_belief_field.py │ ├── demo_ch3_object_graph.py │ ├── demo_ch4_mixture.py │ ├── demo_ch5_viewpoint.py │ └── demo_ch6_control.py │ ├── app.py # Main Gradio app (combines all demos) └── requirements.txt ``` ### Key Implementation Decisions **Simplification 1: Use ground truth segmentation masks from MuJoCo.** MuJoCo gives you per-object segmentation for free (via `mj_render` with segmentation flag). In the minimal version, don't build a learned segmentor — just use GT masks. This lets you focus on the field + graph + dynamics, which is the novel part. State explicitly: "We use oracle segmentation; replacing this with learned segmentation (e.g., SAM2) is future work." **Simplification 2: Train SDF field with direct supervision.** MuJoCo gives you exact object meshes → exact SDFs. Compute GT SDF on a grid, train the field MLP to match. No need for differentiable rendering in the minimal version. The field learns to represent geometry + material under object-conditioned querying. **Simplification 3: Simple material estimation loop.** Don't do full Bayesian inference. Start with a point estimate + learned uncertainty: ```python class MaterialEstimator(nn.Module): def forward(self, interaction_history): """ interaction_history: list of (action, observed_outcome) pairs Returns: material parameter estimates with uncertainty """ # Encode history with a small transformer or RNN h = self.history_encoder(interaction_history) mu = self.mu_head(h) # point estimate sigma = self.sigma_head(h) # uncertainty (softplus to keep positive) return mu, sigma ``` This is simpler than the full variational approach but demonstrates the same concept: uncertainty decreases with more interactions, estimates converge to ground truth. **Simplification 4: Language prior as lookup table.** ```python MATERIAL_PRIORS = { "rubber": {"friction": (0.8, 0.1), "stiffness": (0.01, 0.005), "density": (1100, 200)}, "steel": {"friction": (0.2, 0.05), "stiffness": (200, 20), "density": (7800, 500)}, "wood": {"friction": (0.4, 0.1), "stiffness": (10, 3), "density": (600, 100)}, "glass": {"friction": (0.15, 0.05), "stiffness": (70, 10), "density": (2500, 300)}, "plastic": {"friction": (0.35, 0.08), "stiffness": (3, 1), "density": (1200, 200)}, } def get_language_prior(text): """Returns (mean, std) for each material parameter.""" text = text.lower() for material, params in MATERIAL_PRIORS.items(): if material in text: return params return DEFAULT_PRIOR # uninformative ``` No language model needed. Demonstrates the interface. You state: "In the full system, this lookup would be replaced by a language encoder mapping arbitrary text to material prior distributions." --- ## Data Collection Protocol ### Per-Scene Rollout Collection ```python def collect_rollouts(env, n_rollouts=1000, max_steps=100): """Collect training data from MuJoCo environment.""" dataset = [] for i in range(n_rollouts): obs = env.reset(randomize=True) # Random object positions, orientations trajectory = [] for t in range(max_steps): # Random exploration action (or scripted interactions) action = sample_action(env) # push, poke, drop, slide # Record everything step_data = { 'rgb': env.render(camera='front', mode='rgb'), 'depth': env.render(camera='front', mode='depth'), 'rgb_top': env.render(camera='top', mode='rgb'), 'segmentation': env.render(camera='front', mode='segmentation'), # Ground truth from simulator 'gt_poses': env.get_object_poses(), # SE(3) per object 'gt_velocities': env.get_object_velocities(), # 6D per object 'gt_contacts': env.get_contacts(), # contact pairs + forces 'gt_sdf_samples': env.sample_sdf(n=1000), # random 3D points + SDF values 'gt_materials': env.get_material_params(), # per-object material vectors 'action': action, } trajectory.append(step_data) obs = env.step(action) dataset.append(trajectory) return dataset ``` ### Interaction Types for Material Estimation ```python INTERACTION_SCRIPTS = { 'push': lambda env, obj: env.apply_force(obj, direction='x', magnitude=5.0), 'poke': lambda env, obj: env.apply_impulse(obj, direction='z', magnitude=2.0), 'slide': lambda env, obj: env.apply_force(obj, direction='x', magnitude=1.0, duration=0.5), 'drop': lambda env, obj: env.set_position(obj, height=0.3), # drop from 30cm 'tap': lambda env, obj: env.apply_impulse(obj, direction='random', magnitude=0.5), } ``` Each interaction type provides different material information: - Push/slide → friction (how far it goes) - Drop/tap → restitution (how high it bounces) + density (acceleration under gravity) - Poke → stiffness (deformation response; mainly for soft objects) --- ## Training Schedule (Single GPU) | Stage | What | Training Time | Data | |---|---|---|---| | 1 | SDF field (geometry only) | 2-4 hours | 1000 static scenes, GT SDF samples | | 2 | Material estimator | 4-8 hours | 1000 interaction rollouts with GT materials | | 3 | Gating network | 2-4 hours | Same static scenes, multi-object | | 4 | Encoder (vision → state) | 8-16 hours | 1000 rollouts, RGB+depth → state | | 5 | Dynamics model | 8-16 hours | 1000 rollouts, state → next state | | 6 | End-to-end fine-tuning | 16-24 hours | All data, all losses | **Total: ~2-4 days on a single A100 or equivalent.** For V1 / MVP, you can skip stages 4-6 and just show the field + material estimation with GT encoder (use MuJoCo's GT poses/segmentation as the encoder). This is ~1 day of training and still demonstrates the core ideas.