Skip to content

Examples

Example scripts are in the examples/ directory. Each is self-contained and runnable.

Script Description Source
01_quickstart.py Compile full world model, inspect fields and topology
02_ceo_company_model.py CEO use case: model a company in macro/market context
03_government_policy.py Government use case: policy impact analysis
04_computer_use_agent.py Agent use case: user psychology + world context
05_train_financial.py Train on real FRED + Yahoo Finance data
06_curriculum_training.py Full DAG curriculum training
08_corporate_envs.py Corporate multi-agent RL environments
09_robot_envs.py Robot morphology envs from physical dynamics
10_worldenv_trl_training.py WorldEnv + TRL GRPO training (OpenEnv hackathon)

Quick examples

Compile and inspect the full world

from canvas_engineering import compile_schema, ConnectivityPolicy
from general_unified_world_model import World

world = World()
bound = compile_schema(
    world, T=1, H=128, W=128, d_model=64,
    connectivity=ConnectivityPolicy(intra="dense", parent_child="hub_spoke"),
)

print(f"{len(bound.field_names)} fields")
print(f"{bound.layout.num_positions} positions")
print(f"{len(bound.topology.connections)} connections")

Full world canvas

Output: 857 fields on a 128x128 canvas with 11,735 attention connections.

Source: examples/01_quickstart.py

Project for a hedge fund

from general_unified_world_model import World, project
from general_unified_world_model.schema.business import Business

bound = project(
    World(),
    include=[
        "financial",
        "country_us.macro",
        "regime",
        "forecasts.macro",
        "forecasts.financial",
    ],
    entities={"firm_AAPL": Business(), "firm_NVDA": Business()},
    d_model=64,  # H, W auto-sized
)

Hedge fund canvas

Hedge fund projection: financial, macro, regime, forecasts, and two firms on a 64x64 canvas.

Train on heterogeneous data

from general_unified_world_model import (
    World, project, build_world_model,
    FieldEncoder, FieldDecoder, MaskedCanvasTrainer,
    DatasetSpec, InputSpec, OutputSpec, build_mixed_dataloader,
)

bound = project(World(), include=["financial", "country_us.macro", "regime"], d_model=64)

backbone = build_world_model(bound, n_layers=4, n_heads=4, d_ff=256, n_loops=3)
encoder = FieldEncoder(bound)
decoder = FieldDecoder(bound)

macro_spec = DatasetSpec(
    name="FRED",
    input_specs=[
        InputSpec("gdp", "US GDP growth rate", "country_us.macro.output.gdp_nowcast"),
        InputSpec("cpi", "US CPI YoY", "country_us.macro.inflation.headline_cpi"),
    ],
    output_specs=[
        OutputSpec("gdp", "US GDP growth rate", "country_us.macro.output.gdp_nowcast"),
        OutputSpec("cpi", "US CPI YoY", "country_us.macro.inflation.headline_cpi"),
    ],
)
market_spec = DatasetSpec(
    name="Yahoo",
    input_specs=[
        InputSpec("vix", "CBOE VIX index", "financial.equities.vix"),
        InputSpec("ust10y", "10-year Treasury yield", "financial.yield_curves.ten_year"),
    ],
    output_specs=[
        OutputSpec("vix", "CBOE VIX index", "financial.equities.vix"),
        OutputSpec("ust10y", "10-year Treasury yield", "financial.yield_curves.ten_year"),
    ],
)

loader = build_mixed_dataloader(
    bound,
    sources=[(macro_spec, macro_data), (market_spec, market_data)],
    batch_size=32,
)

Source: examples/05_train_financial.py

DAG curriculum with CogVideoX

from general_unified_world_model import DAGCurriculumTrainer

trainer = DAGCurriculumTrainer(
    nodes=dag,
    data_sources=data_sources,
    backbone="cogvideox",
    device="cuda",
)
trainer.run()

Source: examples/06_curriculum_training.py

Inference with GeneralUnifiedWorldModel

from general_unified_world_model import GeneralUnifiedWorldModel

# No checkpoint needed for schema compilation and observation
model = GeneralUnifiedWorldModel(
    include=["financial", "country_us.macro", "regime", "forecasts"],
    d_model=64,
)
model.observe("financial.yield_curves.ten_year", 4.25)
model.observe("country_us.macro.inflation.headline_cpi", 3.1)

# Get canvas state directly
canvas = model.get_canvas()

# Or run diffusion inference (requires trained weights)
# predictions = model.predict(n_steps=50)

Coming soon

Pretrained checkpoints for inference are not yet released. Training is in progress on H100 GPUs. Check back for checkpoint download links.

LLM-powered projection

from general_unified_world_model import llm_project

result = llm_project(
    "I'm a hedge fund PM. Model US macro, rates, credit, "
    "Apple and NVIDIA. I care about recession risk.",
    provider="anthropic",
)
bound = result.compile(T=1, d_model=64)

LLM-driven curriculum design

from general_unified_world_model import build_curriculum, DatasetProfile

curriculum = build_curriculum(
    goal="Fine-tune to learn cardiovascular patient health",
    datasets=[
        DatasetProfile(name="Hospital EHR", description="Electronic health records"),
        DatasetProfile(name="Insurance Claims", description="Claims and diagnoses"),
    ],
)

nodes = curriculum.to_training_nodes()

Visualizations

All visualizations are generated by scripts/generate_assets.py. To regenerate:

python scripts/generate_assets.py

Regime dashboard

Regime state dashboard

Social graph

CEO entity network