Controllers and Policies#

TorchWM provides several controller and policy classes that convert latent representations into environment actions. The right choice depends on whether you plan online (CEM at every step), use a trained actor-critic, or evolve a controller with black-box optimization.

Overview#

Class

Approach

Used by

RSSMPolicy

Cross-entropy method (CEM) planning in latent space

PlaNet, RSSM training

RolloutGenerator

Collects environment rollouts with optional policy

PlaNet, RSSM training

Controller

Learned linear mapping [z, h] → action via CMA-ES

World Models (Ha & Schmidhuber)

IRISActor

CNN + LSTM + action head (trained with REINFORCE)

IRIS

IRISCritic

CNN + LSTM + value head (λ-return baseline)

IRIS

IRISPolicy

Convenience wrapper around IRISActor

IRIS

CNNFeatureExtractor

Shared 4-layer CNN backbone

IRISActor, IRISCritic

RSSMPolicy — CEM planning#

RSSMPolicy runs a cross-entropy method planner inside an RSSM latent-dynamics model. At each environment step it samples candidate action sequences, rolls them out through the RSSM prior, scores them by predicted reward, and refits a Gaussian to the best candidates.

from torchwm import RSSMPolicy

policy = RSSMPolicy(
    model=rssm,               # RecurrentStateSpaceModel instance
    planning_horizon=20,      # H — imagined steps per candidate
    num_candidates=1000,      # N — candidates per iteration
    num_iterations=10,        # I — CEM refitting iterations
    top_candidates=100,       # K — elite candidates kept
    device=torch.device("cuda"),
)

# At each environment step:
action = policy.poll(observation)

The CEM loop:

  1. Initialise Gaussian N(μ, σ) over action sequences of length H.

  2. Sample N candidate sequences.

  3. Roll out each candidate through deterministic_state_fwd + state_prior, scoring by cumulative pred_reward.

  4. Keep the top K elite sequences and refit (μ, σ).

  5. Repeat for I iterations, then execute the first action of the best mean.

Call policy.reset() to zero the hidden and latent state at the start of a new episode.

RolloutGenerator — collecting rollouts#

RolloutGenerator wraps an environment and an optional policy for collecting episode data. It produces Episode objects compatible with PlaNet’s episode-based Memory.

from torchwm import RSSMPolicy, RolloutGenerator, Episode

generator = RolloutGenerator(
    env,                          # Gymnasium-compatible environment
    device=torch.device("cuda"),
    policy=policy,                # RSSMPolicy (or None for random actions)
    episode_gen=lambda: Episode(postprocess_fn),
    max_episode_steps=1000,
    enable_streaming_video=False,  # set True to stream rollouts to disk
)

Collecting data#

# Single random episode (warmup):
episode = generator.rollout_once(random_policy=True)

# Multiple episodes:
episodes = generator.rollout_n(n=5, random_policy=True)

# Single episode with learned policy:
episode = generator.rollout_once(explore=True)

Evaluation#

episode, frames, metrics = generator.rollout_eval()

# metrics contains:
#   "eval/episode_reward"       — total undiscounted return
#   "eval/reconstruction_loss"  — MSE between reconstructed and true frames
#   "eval/reward_pred_loss"     — MSE between predicted and true reward

The evaluation rollout collects reconstructed frames through the RSSM decoder and logs prediction quality metrics.

Streaming video#

generator = RolloutGenerator(
    env, device,
    policy=policy,
    enable_streaming_video=True,
    streaming_video_path="rollouts/",
    streaming_video_fps=20,
)

IRIS Actor-Critic#

IRIS trains a separate actor and critic inside imagined rollouts (no online planning). Both share the same CNN + LSTM architecture but have different output heads.

CNNFeatureExtractor#

The shared 4-layer CNN backbone:

from torchwm import CNNFeatureExtractor

cnn = CNNFeatureExtractor(
    frame_shape=(3, 64, 64),
    output_size=512,
)

frames = torch.randn(4, 3, 64, 64)
features = cnn(frames)   # (4, 512)

Architecture: Conv2D(3→32) → Conv2D(32→64) → Conv2D(64→128) → Conv2D(128→256) → Linear(4096 → output_size), each conv with kernel 3, stride 2, padding 1 and ReLU.

IRISActor#

from torchwm import IRISActor

actor = IRISActor(
    action_size=6,
    hidden_size=512,
    num_layers=4,
    frame_shape=(3, 64, 64),
)

# Forward with time dimension:
frames = torch.randn(4, 20, 3, 64, 64)   # (B, T, C, H, W)
logits, hidden = actor(frames)            # logits: (B, T, 6)

# Single-step action selection:
action = actor.get_action(frame, temperature=1.0, deterministic=False)

The actor processes frames through CNNFeatureExtractor, then a 4-layer LSTM, then a linear action head. For burn-in (initialising the LSTM hidden state from context frames), pass burn_in_frames to forward().

IRISCritic#

from torchwm import IRISCritic

critic = IRISCritic(
    hidden_size=512,
    num_layers=4,
    frame_shape=(3, 64, 64),
)

frames = torch.randn(4, 20, 3, 64, 64)
values, hidden = critic(frames)   # values: (B, T)

The critic matches the actor’s CNN + LSTM architecture but outputs a scalar value instead of action logits. It maintains a separate CNNFeatureExtractor instance from the actor.

IRISPolicy#

from torchwm import IRISPolicy

policy = IRISPolicy(action_size=6)

frames = torch.randn(4, 3, 64, 64)
logits = policy(frames)                # (4, 6)
action = policy.act(frames[0])         # scalar action index
hidden = policy.init_hidden(4, "cuda")

IRISPolicy is a convenience wrapper around IRISActor. It does not automatically create a critic — instantiate IRISCritic separately if needed.

Putting it together#

import torch
import torchwm

actor = torchwm.IRISActor(action_size=6)
critic = torchwm.IRISCritic()

# Imagined rollout loop:
frames = torch.randn(1, 20, 3, 64, 64)
action_logits, _ = actor(frames)       # (1, 20, 6)
values, _ = critic(frames)             # (1, 20)

# REINFORCE with λ-return baseline:
advantages = ...   # computed from values and predicted rewards
actor_loss = -(action_logits * advantages).mean()
critic_loss = F.mse_loss(values, target_values)

Controller — CMA-ES linear policy#

Controller is a simple linear layer that maps the concatenated latent and deterministic states to an action vector:

from world_models.models.controller import Controller

ctrl = Controller(latent_size=32, hidden_size=256, action_size=3)
action = ctrl(torch.cat([z, h], dim=-1))   # (B, action_size)

Weights are trained with CMA-ES (black-box evolution) rather than gradient descent. See train_controller.py in the World Models (Ha & Schmidhuber) pipeline for a complete example.

Export#

RSSM policies and IRIS actor-critic networks can be exported to ONNX / TorchScript for deployment:

# Export RSSMPolicy's underlying RSSM model:
rssm.export("rssm.onnx", format="onnx", example_inputs=...)

# Export IRIS actor:
actor.export("actor.onnx", format="onnx", example_inputs=...)

See Also#