One Interface Runs Dreamer, IRIS, and DIAMOND: Simulusion, an Open-Source World Model Library

Anyone working on model-based reinforcement learning (MBRL) has probably run into the same headache: switch to a different world model, and you have to rewrite the training scripts, planner, and evaluation pipeline all over again. Dreamer uses RSSM, IRIS uses a token Transformer, and DIAMOND uses diffusion models — their internal structures are wildly different, yet what they need to accomplish externally is highly similar: encode observations, predict the next step, and roll out imagined trajectories. Simulusion, an MIT-licensed PyTorch library, aims to solve exactly this problem: use a unified WorldModel protocol to bring these models under one roof, so a planner written once can run anywhere. The project author posted on the MBRL community on Reddit seeking feedback, and the questions raised happen to hit the key pain points of this kind of abstraction design.

Core Abstraction: The WorldModel Protocol

The design center of Simulusion is a WorldModel protocol consisting of three methods: encode(frames, actions) encodes observation frames and action sequences into latent states; step(state, action) returns rewards and termination signals from the model's own prediction head; rollout(state, actions) unfolds multi-step trajectories in latent space.

The cleverness of this division is that it completely decouples "how the model represents things internally" from "how the planner uses them." RSSM's recurrent state, Transformer's token sequences, and diffusion models' intermediate denoising states are all compressed into an opaque state object. The planner doesn't need to know whether the state contains vectors or tensor collections — it just calls step and rollout. According to the author, two categories of models are currently integrated into this protocol: the RSSM family (Dreamer, PlaNet) and the token Transformer family (IRIS), with diffusion models (DIAMOND) also within the supported scope.

The author also mentions a capability that naturally emerges from this: world models can be trained independently of agents, using only recorded episode data, and once trained, any agent that follows the protocol can directly use it for planning. This has significant engineering implications — model training and policy learning can be advanced by different teams on different timelines.

RSSM (Recurrent State Space Model) is the core component of the Dreamer series, combining a recurrent neural network with stochastic latent variables: the recurrent part maintains a deterministic hidden state responsible for carrying memory across time steps; the stochastic part uses discrete or continuous latent variables to capture future uncertainty. Training typically combines KL balancing, reconstruction loss, and reward prediction loss. PlaNet and DreamerV1-V3 all iterate on this foundation, with differences mainly in loss design, the number of discrete latent variables, and sequence modeling approaches. Because it has both recurrent memory and stochastic sampling capabilities, the planner can perform imagined rollouts in latent space — which is exactly why Simulusion can encapsulate it with a unified state object.

Token Transformer world models are represented by IRIS, which compresses observations into discrete token sequences via VQ-VAE or similar tokenizers, then uses an autoregressive Transformer to model transitions between tokens. Unlike RSSM's continuous latent space, the state here is naturally a string of tokens, and predicting the next step is equivalent to generating token by token. These models have better interpretability and compatibility with language model architectures, but sequence length and tokenization quality directly limit planning accuracy. Simulusion bringing it and RSSM into the same protocol precisely demonstrates that the state object must remain opaque — the planner cannot depend on its internal structure.

Diffusion world models are represented by DIAMOND, which uses diffusion or denoising processes to generate the next frame observation. Its "latent state" is closer to an intermediate representation in the denoising trajectory rather than a compressed low-dimensional vector. These models excel in visual fidelity, but single-step prediction requires multiple denoising iterations, making planning computationally much more expensive than RSSM. The number of step and rollout calls in the protocol thus becomes a practical performance bottleneck, and integration often requires trade-offs between sampling steps and planning depth.

Supporting Toolchain: From Latent Probes to Environment Adapters

Beyond the protocol itself, Simulusion also packages several components that are repeatedly needed in MBRL experiments.

The latent space probe tools are one interesting part. The library includes three metrics — linear probes, DCI, and MIG — to test whether the latents learned by the model actually encode the true state of the environment. This has long been neglected in MBRL: a model performing well in imagined rollouts doesn't mean its representations are meaningful — it might just be overfitting to the reward signal. Having ready-made probe tools at least allows researchers to track "representation quality" as a measurable metric.

The environment adaptation layer unifies DeepMind Control, Atari, Gymnasium, MuJoCo, and Brax under a single make_env entry point. The API differences among these environments are tedious but real, and eliminating this glue code is genuinely helpful for reproducing experiments. Additionally, new models can be registered via decorators and immediately reuse the same Trainer, CLI, and toolchain, lowering the integration cost.

Linear probes, DCI, and MIG are three commonly used representation quality metrics. Linear probes train a linear classifier on frozen latents to predict true state factors — high accuracy indicates the latent is linearly separable for that factor. DCI (Disentanglement, Completeness, Informativeness) scores across three dimensions — disentanglement, completeness, and informativeness — requiring a regressor from latents to true factors to be fitted first. MIG (Mutual Information Gap) is based on mutual information, measuring whether each true factor is primarily encoded by a single latent dimension. The three have different emphases: linear probes look at predictability, DCI looks at disentangled structure, and MIG looks at information concentration. Using them together in MBRL can distinguish between "representations are genuinely good" and "just overfitting to reward."

Two Questions from the Author — Also the Soft Spots of This Abstraction

The author doesn't shy away from design uncertainties in the post, and the first question is particularly worth expanding on: is the encode/step/decode/rollout interface general enough?

He specifically names TD-MPC-style models as a potential counterexample — these models don't decode during planning. Simulusion's protocol includes a decode step, and if a model's workflow simply doesn't need to reconstruct observations from latent states, decode becomes semantically awkward: either force-implement a method that isn't used, or mark it as optional in the protocol — the latter weakening the protocol's uniformity. This is actually a trade-off that all "unified interface" designs must face: the higher the abstraction level, the more likely it is to leave ill-fitting edges on specific models.

The second question is roadmap priority: should DreamerV3 or TD-MPC2 be integrated next? From the community discussion perspective, this choice itself reflects the library's positioning dilemma — DreamerV3 represents the latest mature form of the RSSM approach, and integrating it is more about deepening existing capabilities; TD-MPC2 directly touches on the decode controversy mentioned above, and integrating it is more like a stress test of the protocol's boundaries.

TD-MPC-style models have a core philosophy of being "decoder-free": they don't reconstruct observations, but directly predict rewards, values, and terminal signals from latent states, and use model predictive control (MPC) with learned latent dynamics for short-horizon optimization during planning. Since the entire pipeline doesn't need to restore latents to pixels, the decode step is redundant or even harmful for them — force-implementing it would introduce meaningless computation and training objectives. This contrasts sharply with the Dreamer family, which must constrain representations through reconstruction or contrastive learning. If Simulusion wants to be compatible with TD-MPC2, it must clarify whether decode is an optional capability or a protocol requirement — this directly determines whether the protocol remains unified or moves toward modularization.

The Value and Risks of This Kind of Library

Unified interfaces are genuinely needed in the MBRL community. Implementations in academia and industry are severely fragmented, and the same Dreamer can have subtle but reproducibility-affecting differences across different codebases. A neutral, protocol-driven library that can maintain abstraction consistency is valuable for comparative research.

The risks are equally obvious. Once the abstraction layer is finalized, integrating subsequent models may become a matter of forcing square pegs into round holes; if the protocol keeps adding optional methods for compatibility, it will eventually degrade into "can hold anything, guarantees nothing." The author's approach of openly seeking feedback and proactively exposing design questions is much healthier than working in isolation, but whether the protocol ultimately holds up depends on whether the next few models integrated truly don't need patches.

For those working on MBRL, the point worth watching about this project isn't how many models it currently supports, but that it raises a testable question: what should the common interface for world models actually look like? The answer to this question will directly affect the cost of model reuse and comparative research in the future.

Project address: https://github.com/paramthakkar123/simulusion , website: https://simulusion.vercel.app/ .

One-line Summary

Simulusion uses a unified WorldModel protocol to bring heterogeneous world models under one roof, letting planners be written once and reused across models.

Paragraph Summary

Simulusion is an MIT-licensed PyTorch library that attempts to solve the fragmentation of world model interfaces in model-based reinforcement learning with a unified WorldModel protocol. The protocol consists of three methods — encode, step, and rollout — encapsulating the internal states of RSSM, token Transformers, and diffusion models into opaque state objects, allowing planners to be reused across models without concern for internal structure. The library also provides latent space quality metrics such as linear probes, DCI, and MIG, as well as a make_env entry point unifying DeepMind Control, Atari, Gymnasium, MuJoCo, and Brax. The author openly points out that the protocol may not fit TD-MPC-style models that don't decode, exposing the tension between unified abstraction and model specificity, and poses the roadmap question of whether to integrate DreamerV3 or TD-MPC2 next to the community.