Skip to content

Learned diffusion: replace or gate voxel nutrient diffusion with a model #79

Description

@biosynthart

Goal

Replace or approximate the O(N_cells × 4) lateral nutrient diffusion step (Phase 4 of engine tick loop) with a learned forward pass. This is the first simulation phase to be replaced by a BYOM adapter, establishing the pattern for future phase replacements.

Diffusion was disabled (nutrient_diffusion_enabled: false) because it was expensive. A learned model makes it cheap enough to always enable — or use as an adaptive gate that skips physics diffusion when gradients are flat.

Background

Current diffusion (in world_processes.py):

  • Runs every DIFFUSION_PERIOD=5 ticks
  • For each cell in nutrients_fast and nutrients_slow layers: 4 neighbor lookups, blend computation
  • O(N_cells × 4) dict operations per layer per period
  • Gated by NUTRIENT_DIFFUSION_ENABLED toggle

Diffusion is a smooth local operator — exactly the kind of thing convolutions learn naturally. The model predicts the delta (small perturbation), not the full next state.

See: docs/ECOSIM_PARAMETER_TELEMETRY_SPACE.md § "Concrete Use Case B"

Two Modes

Mode 1: Gating (Phase A — lower risk)

Train a lightweight predictor that estimates |delta| for each layer. If below ε, skip the physics diffusion step entirely. This gives adaptive frequency without replacing the physics.

  • Inputs: voxel grid snapshot at tick t (5 layers, downsampled to 8³ or 16³)
  • Targets: binary — "diffusion delta > ε in next period?"
  • Model: small MLP on pooled features, or tiny CNN
  • Benefit: skip diffusion during steady-state, run it after rain/decomposition blooms

Mode 2: Replacement (Phase B — full replacement)

Train a model that predicts the diffusion delta directly.

  • Inputs: voxel grid snapshot at tick t (5 layers as channels, 32³ spatial)
  • Targets: voxel grid delta over one diffusion period (t → t+5), per layer
  • Model: U-Net or small 3D CNN; predict delta since diffusion is a small perturbation
  • Benefit: O(1) forward pass vs. O(N_cells × 4) neighbor lookups

Adapter Protocol

Extend the BYOM pattern from model_adapter.py:

@runtime_checkable
class VoxelAdapter(Protocol):
    """Voxel-level model adapter for learned diffusion."""

    def context_spec(self) -> ContextSpec:
        """Declare input features (voxel layers + metadata)."""
        ...

    def infer_diffusion_delta(
        self,
        voxel_grid: list[list[float]],  # flattened or structured
        biome_config: dict[str, float],
    ) -> list[list[float]]:
        """Predict diffusion delta for each layer."""
        ...

    def should_skip(self, ...) -> bool:
        """Gating mode: predict whether physics diffusion can be skipped."""
        ...

Training Data Collection

  1. Enable nutrient_diffusion_enabled: true in sim_config
  2. Run 500–1000 sims across all biomes with varied configs (from Surrogate model + sensitivity analysis over sim_config parameter space #78 batch runner)
  3. Log voxel snapshots every 5 ticks (one diffusion period) via telemetry (Telemetry emitter: config snapshot + time-series aggregates + event batching #77)
  4. Compute delta = snapshot[t+5] - snapshot[t] per layer

Module Structure

server/ecosim/adapters/voxel.py     # VoxelAdapter protocol + LearnedDiffusionAdapter
scripts/train_diffusion_model.py    # training script for U-Net/CNN

Design Constraints

  • Follows existing BYOM pattern: context spec → flat float vectors → predictions
  • Model weights saved as JSON (consistent with weights/motion_v0.json)
  • Gating mode works without PyTorch — simple threshold on pooled features
  • Replacement mode needs numpy at minimum, optionally PyTorch for training

Acceptance Criteria

Phase A (Gating)

  • VoxelAdapter protocol defined in model_adapter.py or new module
  • Gating model trained and saved as JSON weights
  • Engine skips diffusion when gating model predicts |delta| < ε
  • Benchmark: same output as always-on diffusion, fewer diffusion steps executed

Phase B (Replacement)

  • U-Net/CNN trained on voxel snapshot → delta data
  • LearnedDiffusionAdapter plugs into engine Phase 4
  • Accuracy: learned delta matches physics within tolerance (<5% MSE per layer)
  • Speedup: forward pass faster than O(N_cells × 4) neighbor lookups

Related

Metadata

Metadata

Assignees

No one assigned

    Labels

    shelvedDeferred until after scalability work

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions