Spaces:

Nirav-Madhani
/

WorldRepresentation

Running

App Files Files Community

WorldRepresentation / docs /06_mujoco_scenes_and_code.md

Nirav Madhani

Import World Representation Space

5af572e 4 months ago

preview code

Raw

History Blame Contribute Delete

15 kB

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

<mujoco model="indistinguishable_pair">
  <worldbody>
    <body name="table">
      <geom type="box" size="0.5 0.5 0.02" pos="0 0 0.5" rgba="0.6 0.6 0.6 1"/>
    </body>
    
    <!-- Heavy cube: steel, 2kg -->
    <body name="cube_heavy" pos="-0.15 0 0.55">
      <joint type="free"/>
      <geom type="box" size="0.03 0.03 0.03" mass="2.0" rgba="0.8 0.2 0.2 1"
            friction="0.5 0.005 0.001"/>
    </body>
    
    <!-- Light cube: foam, 0.02kg -->
    <body name="cube_light" pos="0.15 0 0.55">
      <joint type="free"/>
      <geom type="box" size="0.03 0.03 0.03" mass="0.02" rgba="0.8 0.2 0.2 1"
            friction="0.3 0.005 0.001"/>
    </body>
    
    <!-- Pusher (robot finger proxy) -->
    <body name="pusher" pos="0 -0.15 0.55">
      <joint type="slide" axis="0 1 0" range="-0.3 0.3"/>
      <joint type="slide" axis="1 0 0" range="-0.3 0.3"/>
      <geom type="sphere" size="0.015" rgba="0.3 0.3 0.8 1"/>
    </body>
  </worldbody>
  
  <actuator>
    <velocity joint="pusher_slide_y" kv="50"/>
    <velocity joint="pusher_slide_x" kv="50"/>
  </actuator>
</mujoco>

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:

<visual>
  <!-- Top-down -->
  <camera name="top" pos="0 0 1.5" quat="1 0 0 0"/>
  <!-- Front -->
  <camera name="front" pos="0 -0.8 0.8" xyaxes="1 0 0 0 0.707 0.707"/>
  <!-- Side -->
  <camera name="side" pos="0.8 0 0.8" xyaxes="0 1 0 -0.707 0 0.707"/>
  <!-- Ego (wrist-like) -->
  <camera name="ego" pos="0 -0.3 0.65" xyaxes="1 0 0 0 0.5 0.866"/>
</visual>

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:

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.

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

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

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.