Causal World Models for Physical Reasoning
Project description
PhysiCausal
Causal World Models for Physical Reasoning
Installation • Quickstart • Features • Model Zoo • Documentation • Citation
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 buffertrain()— One-line training function with built-in causal validation checkpointsReplayBuffer— 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
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
2b5e14b0d1991f0e38d0e4c7a3ff5b33c549b5da7b8ae90fa7cc884e1f80a4a7
|
|
| MD5 |
0e47e58d39ec048e1c8bbdb268a83e73
|
|
| BLAKE2b-256 |
69882d0bc8b2a96d918ae39c660055f4b2bb16dd06ff57be0e0885d4621f0e59
|
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
447b556f21d5f589302746753339fa58391a5331817f1263a2cae5d0c18f940f
|
|
| MD5 |
62ce0ce5dc09641d516786596ee15e73
|
|
| BLAKE2b-256 |
ff7f57b9c5de169acc16521dd4a7f8e194390f5a1d50eed4045bae05aac06ec4
|