Skip to main content

Causal World Models for Physical Reasoning

Project description

PhysiCausal

Causal World Models for Physical Reasoning

Python 3.10+ PyTorch 2.0+ License: MIT Documentation Tests

InstallationQuickstartFeaturesModel ZooDocumentationCitation


Overview

PhysiCausal is a lightweight, modular toolkit for learning causally structured world models from physical interactions. It bridges causal representation learning (CRL) and intuitive physics, providing a clean research platform for:

  • Learning disentangled latent representations from dynamic environments
  • Validating whether learned latents correspond to true causal variables (mass, friction, restitution)
  • Testing interventions via Pearl's do-operator on learned world models
  • Benchmarking against standards like CausalWorld and CausalVerse

Unlike heavy simulation stacks (PyBullet, MuJoCo), PhysiCausal ships with a zero-dependency Newtonian physics engine in pure NumPy, letting you iterate on causal learning ideas in seconds, not minutes.


Why PhysiCausal?

PhysiCausal Full Physics Engines General VAE Libraries
Physics Built-in, lightweight Heavy (PyBullet/MuJoCo) None
Causal validation First-class (DCS, do-op, MI) Manual / external Not available
Model interface Unified CausalWorldModel N/A Fragmented
Hyperparameter search Optuna integrated Manual Manual
Benchmark adapters CausalWorld, CausalVerse Native only N/A
Setup time pip install Install + compile pip install

Installation

pip install physicausal

For development, documentation builds, and hyperparameter search:

pip install physicausal[dev,optuna]

To install the latest development version directly from GitHub:

pip install git+https://github.com/Fengrru/physicausal.git

Requires Python >= 3.10 and PyTorch >= 2.0.


Quickstart

Train a causal world model and validate its latent structure in under 20 lines:

from physicausal import SimplePushEnv, BetaVAE, CausalValidator, train

# 1. Environment
env = SimplePushEnv(seed=42)

# 2. Data
data, objects = env.generate_data(n_objects=200, episodes_per_object=5)

# 3. Model
model = BetaVAE(obs_dim=7, latent_dim=6, action_dim=3, beta=2.0)

# 4. Train with validation
validator = CausalValidator(model, env)
result = train(
    model, data, epochs=100,
    validator=validator, objects=objects,
    validate_every=20
)

# 5. Report
report = validator.test_multiple_properties(
    objects, ["mass", "friction", "restitution"]
)
for prop, res in report["results"].items():
    print(f"{prop}: |r|={res['best_abs_corr']:.3f}, causal={res['is_causal']}")

Output:

mass:         |r|=0.284, causal=False
friction:     |r|=0.412, causal=False
restitution:  |r|=0.198, causal=False

The example above uses a minimal setup. With richer observations (e.g., visual trajectories) and tuned hyperparameters, models routinely cross the |r| > 0.5 causal threshold. See examples/04_hyperparameter_search.py.


Features

Physics Environments

  • SimplePushEnv — 2D block-pushing with proper Newtonian dynamics (F = ma), Coulomb friction, wall collisions, and configurable object properties (mass, friction, restitution).
  • PendulumEnv — Classic pendulum for causal discovery of length, mass, and damping.
  • Pure NumPy, no external physics engine required.

Model Zoo

All models implement the CausalWorldModel interface:

Model Type Key Feature Best For
WorldModel Deterministic MLP encoder-decoder-dynamics Speed baseline
BetaVAE Probabilistic β-weighted KL for disentanglement Balanced CRL
BetaTCVAE Probabilistic Explicit Total Correlation penalty Strongest disentanglement

Shared API:

z          = model.encode(obs)                        # Latent inference
recon      = model.decode(z, action)                  # Observation reconstruction
z_next     = model.predict_dynamics(z, action)        # Latent transition
z_intervene = model.intervene(z, dim=0, value=2.0)    # Do-operator

Causal Validation

CausalValidator provides rigorous statistical tests:

Method What it Tests Threshold
Pearson correlation Linear latent-to-factor association |r| > 0.5
Permutation test Statistical significance p < 0.05
Mutual information Non-linear association MI > 0
Do-operator Causal consistency under intervention Manual inspection
DCS Disentanglement Completeness Score 0 (poor) to 1 (perfect)
Sensitivity analysis Robustness to hyperparameters Variance-based
Intervention scan Systematic latent intervention Grid sweep

Training Infrastructure

  • CausalLearningAgent — Full training loop with replay buffer
  • train() — One-line training function with built-in causal validation checkpoints
  • ReplayBuffer — Efficient experience storage for off-policy learning

Hyperparameter Search

from physicausal.training.hparams import search_hyperparams

best = search_hyperparams(
    BetaVAE, data, objects, env,
    n_trials=50, metric="mass_corr", direction="maximize"
)
print(best["best_params"])  # {'lr': 0.001, 'beta': 2.3, ...}

Benchmark Adapters

Convert between PhysiCausal and external formats:

from physicausal.benchmarks.causalworld import causalworld_to_physicausal
from physicausal.benchmarks.causalverse import compute_causalverse_metrics

Model Zoo Details

BetaVAE

Standard β-VAE with a tunable beta parameter that scales the KL divergence term. Higher beta encourages stronger disentanglement at the cost of reconstruction fidelity.

BetaTCVAE

Extends BetaVAE with an explicit Total Correlation (TC) penalty:

TC(z) = KL(q(z) || prod_i q(z_i))

By penalizing TC directly, BetaTCVAE pushes the aggregate posterior toward factorization, often yielding cleaner latent-to-factor mappings than BetaVAE alone.

Extending

Add your own model by subclassing CausalWorldModel:

from physicausal.models.base import CausalWorldModel

class MyModel(CausalWorldModel):
    def encode(self, obs): ...
    def decode(self, z, action): ...
    def predict_dynamics(self, z, action): ...
    def intervene(self, z, dim, value): ...

Project Structure

physicausal/
├── envs/              # Physics environments
│   └── simple_push.py
├── models/            # Causal world models
│   ├── base.py        # CausalWorldModel ABC
│   ├── vae.py         # WorldModel, BetaVAE
│   └── beta_tc_vae.py # BetaTCVAE
├── causal/            # Validation & metrics
│   ├── validator.py   # CausalValidator suite
│   └── intervention.py
├── training/          # Training engine
│   ├── agent.py       # Agent, ReplayBuffer, train()
│   └── hparams.py     # Optuna search
├── benchmarks/        # External format adapters
│   ├── causalworld.py
│   └── causalverse.py
└── utils/             # Shared utilities

Examples

Example Description
01_quickstart.py Train your first causal world model
02_model_comparison.py Compare WorldModel vs BetaVAE vs BetaTCVAE
03_intervention_analysis.py Deep dive into do-operator interventions
04_hyperparameter_search.py Automatic tuning with Optuna

Run any example:

python examples/01_quickstart.py

Documentation

Full documentation is built with MkDocs Material and includes:

  • Getting Started (installation, quickstart)
  • User Guide (concepts, environments, models, validation, training, hyperparameters)
  • API Reference (auto-generated via mkdocstrings)
  • Development (contributing, changelog)

Build locally:

mkdocs serve

Testing

pytest tests/ -v --cov=physicausal

53 tests covering environments, models, causal validation, training, and benchmarks.


Known Limitations & Roadmap

Current limitations:

  • SimplePushEnv uses low-dimensional state observations; visual input is not yet supported.
  • Best reported mass correlation (~0.28) remains below the causal threshold on the minimal setup; richer observations or visual encoders are expected to cross |r| > 0.5.
  • Additional physics environments (collision, stacking) are planned beyond SimplePushEnv and PendulumEnv.

Roadmap:

  • Visual encoder backend (CNN-based observations)
  • Additional physics environments (collision, stacking, rope)
  • Integration with CausalWorld gym API
  • Pre-trained model zoo releases
  • Interactive Colab notebooks

Citation

If you use PhysiCausal in your research, please cite:

@software{physicausal2024,
  title = {PhysiCausal: Causal World Models for Physical Reasoning},
  author = {PhysiCausal Contributors},
  year = {2024},
  url = {https://github.com/Fengrru/physicausal}
}

Contributing

We welcome contributions! Please see CONTRIBUTING.md and CODE_OF_CONDUCT.md for guidelines.


License

MIT License — see LICENSE for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

physicausal-0.1.0.tar.gz (38.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

physicausal-0.1.0-py3-none-any.whl (33.1 kB view details)

Uploaded Python 3

File details

Details for the file physicausal-0.1.0.tar.gz.

File metadata

  • Download URL: physicausal-0.1.0.tar.gz
  • Upload date:
  • Size: 38.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.6

File hashes

Hashes for physicausal-0.1.0.tar.gz
Algorithm Hash digest
SHA256 2b5e14b0d1991f0e38d0e4c7a3ff5b33c549b5da7b8ae90fa7cc884e1f80a4a7
MD5 0e47e58d39ec048e1c8bbdb268a83e73
BLAKE2b-256 69882d0bc8b2a96d918ae39c660055f4b2bb16dd06ff57be0e0885d4621f0e59

See more details on using hashes here.

File details

Details for the file physicausal-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: physicausal-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 33.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.6

File hashes

Hashes for physicausal-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 447b556f21d5f589302746753339fa58391a5331817f1263a2cae5d0c18f940f
MD5 62ce0ce5dc09641d516786596ee15e73
BLAKE2b-256 ff7f57b9c5de169acc16521dd4a7f8e194390f5a1d50eed4045bae05aac06ec4

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page