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

	```xml
	<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:
	```xml
	<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:

	```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.