dreamer4 0.0.31__py3-none-any.whl → 0.1.16__py3-none-any.whl

dreamer4/dreamer4.py

@@ -3,16 +3,18 @@ from __future__ import annotations
 import math
 from math import ceil, log2
 from random import random
+from contextlib import nullcontext
 from collections import namedtuple
-from functools import partial
-from dataclasses import dataclass
+from functools import partial, wraps
+from dataclasses import dataclass, asdict
 
 import torch
 import torch.nn.functional as F
 from torch.nested import nested_tensor
-from torch.distributions import Normal
+from torch.distributions import Normal, kl
 from torch.nn import Module, ModuleList, Embedding, Parameter, Sequential, Linear, RMSNorm, Identity
-from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
+from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, full, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
+from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 
 import torchvision
 from torchvision.models import VGG16_Weights
@@ -20,11 +22,13 @@ from torchvision.models import VGG16_Weights
 from torch.optim import Optimizer
 from adam_atan2_pytorch import MuonAdamAtan2
 
-from x_mlps_pytorch.normed_mlp import create_mlp
 from x_mlps_pytorch.ensemble import Ensemble
+from x_mlps_pytorch.normed_mlp import create_mlp
 
 from hyper_connections import get_init_and_expand_reduce_stream_functions
 
+from vit_pytorch.vit_with_decorr import DecorrelationLoss
+
 from assoc_scan import AssocScan
 
 # ein related
@@ -35,12 +39,15 @@ from assoc_scan import AssocScan
 # d - feature dimension
 # f - frequencies (rotary)
 # l - logit / predicted bins
+# y - layers of transformer
 # p - positions (3 for spacetime in this work)
 # t - time
 # na - action dimension (number of discrete and continuous actions)
 # g - groups of query heads to key heads (gqa)
 # vc - video channels
 # vh, vw - video height and width
+# mtp - multi token prediction length
+# v - video viewpoints
 
 import einx
 from einx import add, multiply
@@ -63,22 +70,96 @@ except ImportError:
 
 LinearNoBias = partial(Linear, bias = False)
 
-TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips'))
+TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips', 'time_decorr', 'space_decorr'))
+
+WorldModelLosses = namedtuple('WorldModelLosses', ('flow', 'rewards', 'discrete_actions', 'continuous_actions'))
 
-WorldModelLosses = namedtuple('WorldModelLosses', ('flow', 'reward', 'behavior_clone'))
+AttentionIntermediates = namedtuple('AttentionIntermediates', ('next_kv_cache', 'normed_inputs'))
+
+TransformerIntermediates = namedtuple('TransformerIntermediates', ('next_kv_cache', 'normed_time_inputs', 'normed_space_inputs'))
+
+MaybeTensor = Tensor | None
 
 @dataclass
 class Experience:
     latents: Tensor
-    video: Tensor | None = None
+    video: MaybeTensor = None
+    proprio: MaybeTensor = None
+    agent_embed: MaybeTensor = None
     rewards: Tensor | None = None
-    actions: tuple[Tensor, Tensor] | None = None
-    log_probs: tuple[Tensor, Tensor] | None = None
-    values: Tensor | None = None
+    actions: tuple[MaybeTensor, MaybeTensor] | None = None
+    log_probs: tuple[MaybeTensor, MaybeTensor] | None = None
+    old_action_unembeds: tuple[MaybeTensor, MaybeTensor] | None = None
+    values: MaybeTensor = None
     step_size: int | None = None
+    lens: MaybeTensor = None
+    is_truncated: MaybeTensor = None
     agent_index: int = 0
     is_from_world_model: bool = True
 
+    def cpu(self):
+        return self.to(torch.device('cpu'))
+
+    def to(self, device):
+        experience_dict = asdict(self)
+        experience_dict = tree_map(lambda t: t.to(device) if is_tensor(t) else t, experience_dict)
+        return Experience(**experience_dict)
+
+def combine_experiences(
+    exps: list[Experiences]
+) -> Experience:
+
+    assert len(exps) > 0
+
+    # set lens if not there
+
+    for exp in exps:
+        latents = exp.latents
+        batch, time, device = *latents.shape[:2], latents.device
+
+        if not exists(exp.lens):
+            exp.lens = full((batch,), time, device = device)
+
+        if not exists(exp.is_truncated):
+            exp.is_truncated = full((batch,), True, device = device)
+
+    # convert to dictionary
+
+    exps_dict = [asdict(exp) for exp in exps]
+
+    values, tree_specs = zip(*[tree_flatten(exp_dict) for exp_dict in exps_dict])
+
+    tree_spec = first(tree_specs)
+
+    all_field_values = list(zip(*values))
+
+    # an assert to make sure all fields are either all tensors, or a single matching value (for step size, agent index etc) - can change this later
+
+    assert all([
+        all([is_tensor(v) for v in field_values]) or len(set(field_values)) == 1
+        for field_values in all_field_values
+    ])
+
+    concatted = []
+
+    for field_values in all_field_values:
+
+        if is_tensor(first(field_values)):
+
+            field_values = pad_tensors_at_dim_to_max_len(field_values, dims = (1, 2))
+
+            new_field_value = cat(field_values)
+        else:
+            new_field_value = first(list(set(field_values)))
+
+        concatted.append(new_field_value)
+
+    # return experience
+
+    concat_exp_dict = tree_unflatten(concatted, tree_spec)
+
+    return Experience(**concat_exp_dict)
+
 # helpers
 
 def exists(v):
@@ -90,6 +171,9 @@ def default(v, d):
 def first(arr):
     return arr[0]
 
+def xnor(x, y):
+    return not (x ^ y)
+
 def has_at_least_one(*bools):
     return sum([*map(int, bools)]) > 0
 
@@ -105,14 +189,55 @@ def sample_prob(prob):
 def is_power_two(num):
     return log2(num).is_integer()
 
+def maybe(fn):
+    def inner(t, *args, **kwargs):
+        if not exists(t) or not exists(fn):
+            return None
+        return fn(t)
+    return inner
+
 # tensor helpers
 
 def is_empty(t):
     return t.numel() == 0
 
+def lens_to_mask(t, max_len = None):
+    if not exists(max_len):
+        max_len = t.amax().item()
+
+    device = t.device
+    seq = torch.arange(max_len, device = device)
+
+    return einx.less('j, i -> i j', seq, t)
+
+def masked_mean(t, mask = None):
+    if not exists(mask):
+        return t.mean()
+
+    if not mask.any():
+        return t[mask].sum()
+
+    return t[mask].mean()
+
 def log(t, eps = 1e-20):
     return t.clamp(min = eps).log()
 
+def mean_log_var_to_distr(
+    mean_log_var: Tensor
+) -> Normal:
+
+    mean, log_var = mean_log_var.unbind(dim = -1)
+    std = (0.5 * log_var).exp()
+    return Normal(mean, std)
+
+def safe_stack(tensors, dim = 0):
+    tensors = [*filter(exists, tensors)]
+
+    if len(tensors) == 0:
+        return None
+
+    return stack(tensors, dim = dim)
+
 def safe_cat(tensors, dim):
     tensors = [*filter(exists, tensors)]
 
@@ -123,6 +248,15 @@ def safe_cat(tensors, dim):
 
     return cat(tensors, dim = dim)
 
+def safe_squeeze_first(t):
+    if not exists(t):
+        return None
+
+    if t.shape[0] != 1:
+        return t
+
+    return rearrange(t, '1 ... -> ...')
+
 def gumbel_noise(t):
     noise = torch.rand_like(t)
     return -log(-log(noise))
@@ -159,6 +293,27 @@ def pad_at_dim(
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
 
+def pad_to_len(t, target_len, *, dim):
+    curr_len = t.shape[dim]
+
+    if curr_len >= target_len:
+        return t
+
+    return pad_at_dim(t, (0, target_len - curr_len), dim = dim)
+
+def pad_tensors_at_dim_to_max_len(
+    tensors: list[Tensor],
+    dims: tuple[int, ...]
+):
+    for dim in dims:
+        if dim >= first(tensors).ndim:
+            continue
+
+        max_time = max([t.shape[dim] for t in tensors])
+        tensors = [pad_to_len(t, max_time, dim = dim) for t in tensors]
+
+    return tensors
+
 def align_dims_left(t, aligned_to):
     shape = t.shape
     num_right_dims = aligned_to.ndim - t.ndim
@@ -174,8 +329,74 @@ def l2norm(t):
 def softclamp(t, value = 50.):
     return (t / value).tanh() * value
 
+def create_multi_token_prediction_targets(
+    t, # (b t ...)
+    steps_future,
+
+): # (b t-1 steps ...), (b t-1 steps) - targets and the mask, where mask is False for padding
+
+    batch, seq_len, device = *t.shape[:2], t.device
+
+    batch_arange = arange(batch, device = device)
+    seq_arange = arange(seq_len, device = device)
+    steps_arange = arange(steps_future, device = device)
+
+    indices = add('t, steps -> t steps', seq_arange, steps_arange)
+    mask = indices < seq_len
+
+    batch_arange = rearrange(batch_arange, 'b -> b 1 1')
+
+    indices[~mask] = 0
+    mask = repeat(mask, 't steps -> b t steps', b = batch)
+
+    out = t[batch_arange, indices]
+
+    return out, mask
+
 # loss related
 
+class LossNormalizer(Module):
+
+    # the authors mentioned the need for loss normalization in the dynamics transformer
+
+    def __init__(
+        self,
+        num_losses: int,
+        beta = 0.95,
+        eps = 1e-6
+    ):
+        super().__init__()
+        self.register_buffer('exp_avg_sq', torch.ones(num_losses))
+        self.beta = beta
+        self.eps = eps
+
+    def forward(
+        self,
+        losses: Tensor | list[Tensor] | dict[str, Tensor],
+        update_ema = None
+    ):
+        exp_avg_sq, beta = self.exp_avg_sq, self.beta
+        update_ema = default(update_ema, self.training)
+
+        # get the rms value - as mentioned at the end of section 3 in the paper
+
+        rms = exp_avg_sq.sqrt()
+
+        if update_ema:
+            decay = 1. - beta
+
+            # update the ema
+
+            exp_avg_sq.lerp_(losses.detach().square(), decay)
+
+        # then normalize
+
+        assert losses.numel() == rms.numel()
+
+        normed_losses = losses / rms.clamp(min = self.eps)
+
+        return normed_losses
+
 class LPIPSLoss(Module):
     def __init__(
         self,
@@ -340,7 +561,9 @@ class ActionEmbedder(Module):
         num_continuous_actions  = 0,
         continuous_norm_stats: tuple[tuple[float, float], ...] | None = None,
         can_unembed = False,
-        unembed_dim = None
+        unembed_dim = None,
+        num_unembed_preds = 1,
+        squeeze_unembed_preds = True # will auto-squeeze if prediction is just 1
     ):
         super().__init__()
 
@@ -378,11 +601,14 @@ class ActionEmbedder(Module):
 
         self.can_unembed = can_unembed
 
+        self.num_unembed_preds = num_unembed_preds
+        self.squeeze_unembed_preds = squeeze_unembed_preds
+
         if not can_unembed:
             return
 
         unembed_dim = default(unembed_dim, dim)
-        self.discrete_action_unembed = Parameter(torch.randn(total_discrete_actions, unembed_dim) * 1e-2)
+        self.discrete_action_unembed = Parameter(torch.randn(total_discrete_actions, num_unembed_preds, unembed_dim) * 1e-2)
 
         discrete_action_index = arange(total_discrete_actions)
 
@@ -396,13 +622,13 @@ class ActionEmbedder(Module):
 
         self.register_buffer('discrete_action_mask', discrete_action_mask, persistent = False)
 
-        self.continuous_action_unembed = Parameter(torch.randn(num_continuous_actions, unembed_dim, 2) * 1e-2)
+        self.continuous_action_unembed = Parameter(torch.randn(num_continuous_actions, num_unembed_preds, unembed_dim, 2) * 1e-2)
 
     def embed_parameters(self):
         return set([*self.discrete_action_embed.parameters(), *self.continuous_action_embed.parameters()])
 
     def unembed_parameters(self):
-        return set([*self.discrete_action_unembed.parameters(), *self.continuous_action_unembed.parameters()])
+        return set([self.discrete_action_unembed, self.continuous_action_unembed])
 
     @property
     def device(self):
@@ -429,12 +655,26 @@ class ActionEmbedder(Module):
         embeds,                          # (... d)
         discrete_action_types = None,    # (na)
         continuous_action_types = None,  # (na)
-        return_split_discrete = False
+        return_split_discrete = False,
+        pred_head_index: int | Tensor | None = None
 
     ):  # (... discrete_na), (... continuous_na 2)
 
+        device = embeds.device
+
         assert self.can_unembed, 'can only unembed for predicted discrete and continuous actions if `can_unembed = True` is set on init'
 
+        # handle only one prediction head during inference
+
+        if exists(pred_head_index) and isinstance(pred_head_index, int):
+            pred_head_index = tensor(pred_head_index, device = device)
+
+        # if pred_head_index given as a solo int, just assume we want to squeeze out the prediction head dimension
+
+        squeeze_one_pred_head = exists(pred_head_index) and pred_head_index.ndim == 0
+
+        # get action types
+
         discrete_action_types, continuous_action_types = tuple(self.cast_action_types(t) for t in (discrete_action_types, continuous_action_types))
 
         # discrete actions
@@ -450,7 +690,13 @@ class ActionEmbedder(Module):
 
                 discrete_action_unembed = discrete_action_unembed[discrete_action_mask]
 
-            discrete_action_logits = einsum(embeds, discrete_action_unembed, '... d, na d -> ... na')
+            if exists(pred_head_index):
+                discrete_action_unembed = discrete_action_unembed.index_select(1, pred_head_index)
+
+            discrete_action_logits = einsum(embeds, discrete_action_unembed, '... d, na mtp d -> mtp ... na')
+
+            if self.squeeze_unembed_preds or squeeze_one_pred_head:
+                discrete_action_logits = safe_squeeze_first(discrete_action_logits)
 
         # whether to split the discrete action logits by the number of actions per action type
 
@@ -471,7 +717,13 @@ class ActionEmbedder(Module):
             if exists(continuous_action_types):
                 continuous_action_unembed = continuous_action_unembed[continuous_action_types]
 
-            continuous_action_mean_log_var = einsum(embeds, continuous_action_unembed, '... d, na d two -> ... na two')
+            if exists(pred_head_index):
+                continuous_action_unembed = continuous_action_unembed.index_select(1, pred_head_index)
+
+            continuous_action_mean_log_var = einsum(embeds, continuous_action_unembed, '... d, na mtp d two -> mtp ... na two')
+
+            if self.squeeze_unembed_preds or squeeze_one_pred_head:
+                continuous_action_mean_log_var = safe_squeeze_first(continuous_action_mean_log_var)
 
         return discrete_action_logits, continuous_action_mean_log_var
 
@@ -480,9 +732,14 @@ class ActionEmbedder(Module):
         embed,
         discrete_temperature = 1.,
         continuous_temperature = 1.,
+        inverse_norm_continuous = None,
+        pred_head_index: int | Tensor | None = None,
+        squeeze = True,
         **kwargs
     ):
-        discrete_logits, continuous_mean_log_var = self.unembed(embed, return_split_discrete = True, **kwargs)
+        inverse_norm_continuous = default(inverse_norm_continuous, self.continuous_need_norm)
+
+        discrete_logits, continuous_mean_log_var = self.unembed(embed, return_split_discrete = True, pred_head_index = pred_head_index, **kwargs)
 
         sampled_discrete = sampled_continuous = None
 
@@ -500,6 +757,12 @@ class ActionEmbedder(Module):
 
             sampled_continuous = mean + std * torch.randn_like(mean) * continuous_temperature
 
+            # maybe inverse norm
+
+            if inverse_norm_continuous:
+                norm_mean, norm_std = self.continuous_norm_stats.unbind(dim = -1)
+                sampled_continuous = (sampled_continuous * norm_std) + norm_mean
+
         return sampled_discrete, sampled_continuous
 
     def log_probs(
@@ -509,12 +772,13 @@ class ActionEmbedder(Module):
         continuous_targets = None,       # (... na)
         discrete_action_types = None,    # (na)
         continuous_action_types = None,  # (na)
+        pred_head_index: int | Tensor | None = None,
         parallel_discrete_calc = None,
         return_entropies = False
     ):
         parallel_discrete_calc = default(parallel_discrete_calc, exists(discrete_targets) and discrete_targets.shape[-1] > 1)
 
-        discrete_action_logits, continuous_action_mean_log_var = self.unembed(embeds, discrete_action_types = discrete_action_types, continuous_action_types = continuous_action_types, return_split_discrete = True)
+        discrete_action_logits, continuous_action_mean_log_var = self.unembed(embeds, pred_head_index = pred_head_index, discrete_action_types = discrete_action_types, continuous_action_types = continuous_action_types, return_split_discrete = True)
 
         # discrete
 
@@ -600,10 +864,7 @@ class ActionEmbedder(Module):
         continuous_entropies = None
 
         if exists(continuous_targets):
-            mean, log_var = continuous_action_mean_log_var.unbind(dim = -1)
-            std = (0.5 * log_var).exp()
-
-            distr = Normal(mean, std)
+            distr = mean_log_var_to_distr(continuous_action_mean_log_var)
             continuous_log_probs = distr.log_prob(continuous_targets)
 
             if return_entropies:
@@ -618,6 +879,64 @@ class ActionEmbedder(Module):
 
         return log_probs, entropies
 
+    def kl_div(
+        self,
+        src: tuple[MaybeTensor, MaybeTensor],
+        tgt: tuple[MaybeTensor, MaybeTensor],
+        reduce_across_num_actions = True
+    ) -> tuple[MaybeTensor, MaybeTensor]:
+
+        src_discrete, src_continuous = src
+        tgt_discrete, tgt_continuous = tgt
+
+        discrete_kl_div = None
+
+        # split discrete if it is not already (multiple discrete actions)
+
+        if exists(src_discrete):
+
+            discrete_split = self.num_discrete_actions.tolist()
+
+            if is_tensor(src_discrete):
+                src_discrete = src_discrete.split(discrete_split, dim = -1)
+
+            if is_tensor(tgt_discrete):
+                tgt_discrete = tgt_discrete.split(discrete_split, dim = -1)
+
+            discrete_kl_divs = []
+
+            for src_logit, tgt_logit in zip(src_discrete, tgt_discrete):
+
+                src_log_probs = src_logit.log_softmax(dim = -1)
+                tgt_prob = tgt_logit.softmax(dim = -1)
+
+                one_discrete_kl_div = F.kl_div(src_log_probs, tgt_prob, reduction = 'none')
+
+                discrete_kl_divs.append(one_discrete_kl_div.sum(dim = -1))
+
+            discrete_kl_div = stack(discrete_kl_divs, dim = -1)
+
+        # calculate kl divergence for continuous
+
+        continuous_kl_div = None
+
+        if exists(src_continuous):
+            src_normal = mean_log_var_to_distr(src_continuous)
+            tgt_normal = mean_log_var_to_distr(tgt_continuous)
+
+            continuous_kl_div = kl.kl_divergence(src_normal, tgt_normal)
+
+        # maybe reduce
+
+        if reduce_across_num_actions:
+            if exists(discrete_kl_div):
+                discrete_kl_div = discrete_kl_div.sum(dim = -1)
+
+            if exists(continuous_kl_div):
+                continuous_kl_div = continuous_kl_div.sum(dim = -1)
+
+        return discrete_kl_div, continuous_kl_div
+
     def forward(
         self,
         *,
@@ -726,11 +1045,12 @@ class Rotary1D(Module):
 
     def forward(
         self,
-        seq_len
+        seq_len,
+        offset = 0
     ):
         device, dtype = self.inv_freq.device, self.inv_freq.dtype
 
-        t = torch.arange(seq_len, device = device).type(dtype)
+        t = torch.arange(seq_len, device = device).type(dtype) + offset
         freqs = einsum(t, self.inv_freq, 'i, j -> i j')
 
         return cat((freqs, freqs), dim = -1)
@@ -740,7 +1060,18 @@ def apply_rotations(
     rotations, # (h n d) | (n d)
     t          # (b h n d)
 ):
-    heads, dtype = t.shape[1], t.dtype
+
+    heads, seq_len, dtype = *t.shape[1:3], t.dtype
+
+    rotations_seq_len = rotations.shape[-2]
+
+    # handle kv caching with rotations
+
+    if rotations_seq_len > seq_len:
+        rotations = rotations[-seq_len:]
+
+    # precision
+
     t = t.float()
 
     # handle gqa for rotary
@@ -877,10 +1208,11 @@ def special_token_mask(q, k, seq_len, num_tokens, special_attend_only_itself = F
 
 def block_mask_special_tokens_right(
     seq_len,
-    num_tokens
+    num_tokens,
+    special_attend_only_itself = False
 ):
     def inner(b, h, q, k):
-        return special_token_mask(q, k, seq_len, num_tokens)
+        return special_token_mask(q, k, seq_len, num_tokens, special_attend_only_itself)
     return inner
 
 def compose_mask(mask1, mask2):
@@ -957,6 +1289,10 @@ class Attention(Module):
         query_heads = None,
         heads = 8,
         pre_rmsnorm = True,
+        gate_values = True,
+        rmsnorm_query = False, # a paper claims that it is better to just norm only the keys https://openreview.net/forum?id=HkztQWZfl2
+        rmsnorm_key = True,
+        value_residual = True
     ):
         super().__init__()
         self.norm = RMSNorm(dim) if pre_rmsnorm else Identity()
@@ -979,12 +1315,33 @@ class Attention(Module):
         self.to_v = LinearNoBias(dim, dim_kv_inner)
         self.to_out = LinearNoBias(dim_q_inner, dim)
 
+        # alphafold gating per head, for attending to nothing
+
+        self.to_gates = None
+
+        if gate_values:
+            self.to_gates = Sequential(
+                LinearNoBias(dim, query_heads),
+                Rearrange('b n h -> b h n 1'),
+                nn.Sigmoid()
+            )
+
         # stability related
 
-        self.q_heads_rmsnorm = MultiHeadRMSNorm(dim_head, heads = query_heads)
-        self.k_heads_rmsnorm = MultiHeadRMSNorm(dim_head, heads = heads)
+        self.q_heads_rmsnorm = MultiHeadRMSNorm(dim_head, heads = query_heads) if rmsnorm_query else nn.Identity()
+        self.k_heads_rmsnorm = MultiHeadRMSNorm(dim_head, heads = heads) if rmsnorm_key else nn.Identity()
+
+        # value residual
+
+        self.to_learned_value_residual_mix = nn.Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b n h -> b h n 1'),
+            nn.Sigmoid()
+        ) if value_residual else None
 
     def muon_parameters(self):
+        # omit the queries and keys for now given what we learned from kimi 2 paper
+
         return [
             *self.to_v.parameters(),
             *self.to_out.parameters(),
@@ -994,8 +1351,9 @@ class Attention(Module):
         self,
         tokens, # (b n d)
         kv_cache = None,
-        return_kv_cache = False,
+        return_intermediates = False,
         rotary_pos_emb = None,
+        residual_values = None,  # (b n h d)
         attend_fn: Callable | None = None
     ):
         tokens, inverse_packed_batch = pack_one(tokens, '* n d')
@@ -1008,11 +1366,28 @@ class Attention(Module):
 
         q, k, v = map(self.split_heads, (q, k, v))
 
+        # handle maybe value residual
+
+        if exists(residual_values):
+            residual_values = rearrange(residual_values, '... n h d -> (...) h n d')
+
+            assert exists(self.to_learned_value_residual_mix)
+
+            learned_mix = self.to_learned_value_residual_mix(tokens)
+
+            v = v.lerp(residual_values, learned_mix)
+
         # qk rmsnorm
 
         q = self.q_heads_rmsnorm(q)
         k = self.k_heads_rmsnorm(k)
 
+        # rotary
+
+        if exists(rotary_pos_emb):
+            q = apply_rotations(rotary_pos_emb, q)
+            k = apply_rotations(rotary_pos_emb, k)
+
         # caching
 
         if exists(kv_cache):
@@ -1020,18 +1395,18 @@ class Attention(Module):
             k = cat((ck, k), dim = -2)
             v = cat((cv, v), dim = -2)
 
-        # rotary
-
-        if exists(rotary_pos_emb):
-            q = apply_rotations(rotary_pos_emb, q)
-            k = apply_rotations(rotary_pos_emb, k)
-
         # attention
 
         attend_fn = default(attend_fn, naive_attend)
 
         out = attend_fn(q, k, v)
 
+        # gate values
+
+        if exists(self.to_gates):
+            gates = self.to_gates(tokens)
+            out = out * gates
+
         # merge heads
 
         out = self.merge_heads(out)
@@ -1042,10 +1417,10 @@ class Attention(Module):
 
         out = inverse_packed_batch(out)
 
-        if not return_kv_cache:
+        if not return_intermediates:
             return out
 
-        return out, stack((k, v))
+        return out, AttentionIntermediates(stack((k, v)), tokens)
 
 # feedforward
 
@@ -1085,6 +1460,7 @@ class AxialSpaceTimeTransformer(Module):
         self,
         dim,
         depth,
+        attn_heads = 8,
         attn_dim_head = 64,
         attn_softclamp_value = 50.,
         time_block_every = 4,
@@ -1093,9 +1469,12 @@ class AxialSpaceTimeTransformer(Module):
         num_residual_streams = 1,
         num_special_spatial_tokens = 1,
         special_attend_only_itself = False,  # this is set to True for the video tokenizer decoder (latents can only attend to itself while spatial modalities attend to the latents and everything)
-        final_norm = True
+        final_norm = True,
+        value_residual = True,               # https://arxiv.org/abs/2410.17897 - but with learned mixing from OSS
+        rnn_time = False
     ):
         super().__init__()
+        assert depth >= time_block_every, f'depth must be at least {time_block_every}'
 
         # hyper connections
 
@@ -1113,6 +1492,24 @@ class AxialSpaceTimeTransformer(Module):
 
         self.time_rotary = Rotary1D(attn_dim_head)
 
+        # project initial for value residuals
+
+        self.value_residual = value_residual
+
+        if value_residual:
+            dim_inner = attn_dim_head * attn_heads
+
+            self.to_value_residual = nn.Sequential(
+                nn.RMSNorm(dim),
+                nn.Linear(dim, dim_inner, bias = False),
+                Rearrange('... (h d) -> ... h d', h = attn_heads)
+            )
+
+        # a gru layer across time
+
+        self.rnn_time = rnn_time
+        rnn_layers = []
+
         # transformer
 
         layers = []
@@ -1124,17 +1521,24 @@ class AxialSpaceTimeTransformer(Module):
             is_time_block = divisible_by(layer_index, time_block_every)
             is_time.append(is_time_block)
 
-            rearrange_to_attend = Rearrange('b t s d -> b s t d') if is_time_block else Identity()
-            rearrange_from_attend = Rearrange('b s t d -> b t s d') if is_time_block else Identity()
+            rearrange_to_attend = Rearrange('b t s ... -> b s t ...') if is_time_block else Identity()
+            rearrange_from_attend = Rearrange('b s t ... -> b t s ...') if is_time_block else Identity()
 
             layers.append(ModuleList([
                 rearrange_to_attend,
                 rearrange_from_attend,
-                hyper_conn(branch = Attention(dim = dim, dim_head = attn_dim_head, **attn_kwargs)),
+                hyper_conn(branch = Attention(dim = dim, heads = attn_heads, dim_head = attn_dim_head, value_residual = value_residual, **attn_kwargs)),
                 hyper_conn(branch = SwiGLUFeedforward(dim = dim, **ff_kwargs))
             ]))
 
+            rnn_layers.append(ModuleList([
+                nn.RMSNorm(dim),
+                nn.GRU(dim, dim, batch_first = True)
+            ]) if is_time_block and rnn_time else None)
+
         self.layers = ModuleList(layers)
+        self.rnn_layers = ModuleList(rnn_layers)
+
         self.is_time = is_time
 
         # final norm
@@ -1145,17 +1549,31 @@ class AxialSpaceTimeTransformer(Module):
 
         self.num_special_spatial_tokens = num_special_spatial_tokens
 
+    def muon_parameters(self):
+        muon_params = []
+
+        for m in self.modules():
+            if isinstance(m, (Attention, SwiGLUFeedforward)):
+                muon_params.extend(m.muon_parameters())
+
+        return muon_params
+
     def forward(
         self,
-        tokens # (b t s d)
-    ):
+        tokens,                          # (b t s d)
+        kv_cache: Tensor | None = None,  # (y 2 b h t d)
+        return_intermediates = False
+
+    ): # (b t s d) | (y 2 b h t d)
+
         batch, time, space_seq_len, _, device = *tokens.shape, tokens.device
 
         assert tokens.ndim == 4
 
         # attend functions for space and time
 
-        use_flex = exists(flex_attention) and tokens.is_cuda
+        has_kv_cache = exists(kv_cache) 
+        use_flex = exists(flex_attention) and tokens.is_cuda and not has_kv_cache # KV cache shape breaks flex attention TODO: Fix
 
         attend_kwargs = dict(use_flex = use_flex, softclamp_value = self.attn_softclamp_value, special_attend_only_itself = self.special_attend_only_itself, device = device)
 
@@ -1163,37 +1581,120 @@ class AxialSpaceTimeTransformer(Module):
 
         time_attend = get_attend_fn(causal = True, seq_len = time, k_seq_len = time, **attend_kwargs)
 
+        # prepare cache
+
+        time_attn_kv_caches = []
+
+        if has_kv_cache:
+            past_tokens, tokens = tokens[:, :-1], tokens[:, -1:]
+
+            rotary_seq_len = 1
+            rotary_pos_offset = past_tokens.shape[1]
+        else:
+            rotary_seq_len = time
+            rotary_pos_offset = 0
+
+        kv_cache = default(kv_cache, (None,))
+
+        iter_kv_cache = iter(kv_cache)
+
         # rotary
 
-        rotary_pos_emb = self.time_rotary(time)
+        rotary_pos_emb = self.time_rotary(rotary_seq_len, offset = rotary_pos_offset)
+
+        # value residual
+
+        residual_values = None
+
+        if self.value_residual:
+            residual_values = self.to_value_residual(tokens)
+
+        # normed attention inputs
+
+        normed_time_attn_inputs = []
+        normed_space_attn_inputs = []
 
         # attention
 
         tokens = self.expand_streams(tokens)
 
-        for (pre_attn_rearrange, post_attn_rearrange, attn, ff), layer_is_time in zip(self.layers, self.is_time):
+        for (pre_attn_rearrange, post_attn_rearrange, attn, ff), maybe_rnn_modules, layer_is_time in zip(self.layers, self.rnn_layers, self.is_time):
 
             tokens = pre_attn_rearrange(tokens)
 
+            # maybe rnn for time
+
+            if layer_is_time and exists(maybe_rnn_modules):
+                rnn_prenorm, rnn = maybe_rnn_modules
+
+                rnn_input, inverse_pack_time = pack_one(tokens, '* t d')
+
+                rnn_out, rnn_hiddens = rnn(rnn_prenorm(rnn_input)) # todo, handle rnn cache
+
+                rnn_out = inverse_pack_time(rnn_out)
+
+                tokens = rnn_out + tokens
+
             # when is a axial time attention block, should be causal
 
             attend_fn = time_attend if layer_is_time else space_attend
 
             layer_rotary_pos_emb = rotary_pos_emb if layer_is_time else None
 
+            # maybe past kv cache
+
+            maybe_kv_cache = next(iter_kv_cache, None) if layer_is_time else None
+
+            # residual values
+
+            layer_residual_values = maybe(pre_attn_rearrange)(residual_values)
+
             # attention layer
 
-            tokens = attn(tokens, rotary_pos_emb = layer_rotary_pos_emb, attend_fn = attend_fn) + tokens
+            tokens, attn_intermediates = attn(
+                tokens,
+                rotary_pos_emb = layer_rotary_pos_emb,
+                attend_fn = attend_fn,
+                kv_cache = maybe_kv_cache,
+                residual_values = layer_residual_values,
+                return_intermediates = True
+            )
 
             tokens = post_attn_rearrange(tokens)
 
             # feedforward layer
 
-            tokens = ff(tokens) + tokens
+            tokens = ff(tokens)
+
+            # save kv cache if is time layer
+
+            if layer_is_time:
+                time_attn_kv_caches.append(attn_intermediates.next_kv_cache)
+
+            # save time attention inputs for decorr
+
+            space_or_time_inputs = normed_time_attn_inputs if layer_is_time else normed_space_attn_inputs
+
+            space_or_time_inputs.append(attn_intermediates.normed_inputs)
 
         tokens = self.reduce_streams(tokens)
 
-        return self.final_norm(tokens)
+        out = self.final_norm(tokens)
+
+        if has_kv_cache:
+            # just concat the past tokens back on for now, todo - clean up the logic
+            out = cat((past_tokens, out), dim = 1)
+
+        if not return_intermediates:
+            return out
+
+        intermediates = TransformerIntermediates(
+            stack(time_attn_kv_caches),
+            safe_stack(normed_time_attn_inputs),
+            safe_stack(normed_space_attn_inputs)
+        )
+
+        return out, intermediates
 
 # video tokenizer
 
@@ -1219,12 +1720,15 @@ class VideoTokenizer(Module):
         per_image_patch_mask_prob = (0., 0.9), # probability of patch masking appears to be per image probabilities drawn uniformly between 0. and 0.9 - if you are a phd student and think i'm mistakened, please open an issue
         lpips_loss_network: Module | None = None,
         lpips_loss_weight = 0.2,
+        encoder_add_decor_aux_loss = False,
+        decor_auxx_loss_weight = 0.1,
+        decorr_sample_frac = 0.25,
         nd_rotary_kwargs: dict = dict(
             rope_min_freq = 1.,
             rope_max_freq = 10000.,
             rope_p_zero_freqs = 0.
         ),
-        num_residual_streams = 1
+        num_residual_streams = 1,
     ):
         super().__init__()
 
@@ -1305,6 +1809,7 @@ class VideoTokenizer(Module):
             time_block_every = time_block_every,
             num_special_spatial_tokens = num_latent_tokens,
             num_residual_streams = num_residual_streams,
+            special_attend_only_itself = True,
             final_norm = True
         )
 
@@ -1318,10 +1823,24 @@ class VideoTokenizer(Module):
         if self.has_lpips_loss:
             self.lpips = LPIPSLoss(lpips_loss_network)
 
+        # decorr aux loss
+        # https://arxiv.org/abs/2510.14657
+
+        self.encoder_add_decor_aux_loss = encoder_add_decor_aux_loss
+        self.decorr_aux_loss_weight = decor_auxx_loss_weight
+
+        self.decorr_loss = DecorrelationLoss(decorr_sample_frac, soft_validate_num_sampled = True) if encoder_add_decor_aux_loss else None
+
     @property
     def device(self):
         return self.zero.device
 
+    def muon_parameters(self):
+        return [
+            *self.encoder_transformer.muon_parameters(),
+            *self.decoder_transformer.muon_parameters()
+        ]
+
     @torch.no_grad()
     def tokenize(
         self,
@@ -1425,7 +1944,7 @@ class VideoTokenizer(Module):
 
         # encoder attention
 
-        tokens = self.encoder_transformer(tokens)
+        tokens, (_, time_attn_normed_inputs, space_attn_normed_inputs) = self.encoder_transformer(tokens, return_intermediates = True)
 
         # latent bottleneck
 
@@ -1447,19 +1966,30 @@ class VideoTokenizer(Module):
         if self.has_lpips_loss:
             lpips_loss = self.lpips(video, recon_video)
 
+        time_decorr_loss = space_decorr_loss = self.zero
+
+        if self.encoder_add_decor_aux_loss:
+            if exists(time_attn_normed_inputs):
+                time_decorr_loss = self.decorr_loss(time_attn_normed_inputs)
+
+            if exists(space_attn_normed_inputs):
+                space_decorr_loss = self.decorr_loss(space_attn_normed_inputs)
+
         # losses
 
         total_loss = (
             recon_loss +
-            lpips_loss * self.lpips_loss_weight
+            lpips_loss * self.lpips_loss_weight +
+            time_decorr_loss * self.decorr_aux_loss_weight +
+            space_decorr_loss * self.decorr_aux_loss_weight
         )
 
         if not return_all_losses:
             return total_loss
 
-        losses = (recon_loss, lpips_loss)
+        losses = (recon_loss, lpips_loss, decorr_loss)
 
-        return total_loss, TokenizerLosses(losses)
+        return total_loss, TokenizerLosses(*losses)
 
 # dynamics model, axial space-time transformer
 
@@ -1475,13 +2005,15 @@ class DynamicsWorldModel(Module):
         num_latent_tokens = None,
         num_agents = 1,
         num_tasks = 0,
+        num_video_views = 1,
+        dim_proprio = None,
         reward_encoder_kwargs: dict = dict(),
         depth = 4,
         pred_orig_latent = True,   # directly predicting the original x0 data yield better results, rather than velocity (x-space vs v-space)
         time_block_every = 4,      # every 4th block is time
-        attn_kwargs: dict = dict(
-            heads = 8,
-        ),
+        attn_kwargs: dict = dict(),
+        transformer_kwargs: dict = dict(),
+        attn_heads = 8,
         attn_dim_head = 64,
         attn_softclamp_value = 50.,
         ff_kwargs: dict = dict(),
@@ -1493,15 +2025,25 @@ class DynamicsWorldModel(Module):
         num_discrete_actions: int | tuple[int, ...] = 0,
         num_continuous_actions = 0,
         continuous_norm_stats = None,
-        reward_loss_weight = 0.1,
+        multi_token_pred_len = 8,
         value_head_mlp_depth = 3,
         policy_head_mlp_depth = 3,
-        behavior_clone_weight = 0.1,
+        latent_flow_loss_weight = 1.,
+        reward_loss_weight: float | list[float] = 1.,
+        discrete_action_loss_weight: float | list[float] = 1.,
+        continuous_action_loss_weight: float | list[float] = 1.,
         num_latent_genes = 0,                       # for carrying out evolution within the dreams https://web3.arxiv.org/abs/2503.19037
         num_residual_streams = 1,
+        keep_reward_ema_stats = False,
+        reward_ema_decay = 0.998,
+        reward_quantile_filter = (0.05, 0.95),
         gae_discount_factor = 0.997,
         gae_lambda = 0.95,
         ppo_eps_clip = 0.2,
+        pmpo_pos_to_neg_weight = 0.5, # pos and neg equal weight
+        pmpo_reverse_kl = True,
+        pmpo_kl_div_loss_weight = .3,
+        normalize_advantages = None,
         value_clip = 0.4,
         policy_entropy_weight = .01,
         gae_use_accelerated = False
@@ -1535,7 +2077,7 @@ class DynamicsWorldModel(Module):
             )
 
             self.to_latent_pred = Sequential(
-                Reduce('b t n s d -> b t n d', 'mean'),
+                Reduce('b t v n s d -> b t v n d', 'mean'),
                 RMSNorm(dim),
                 LinearNoBias(dim, dim_latent)
             )
@@ -1545,15 +2087,38 @@ class DynamicsWorldModel(Module):
             latent_tokens_to_space = num_latent_tokens // num_spatial_tokens
 
             self.latents_to_spatial_tokens = Sequential(
-                Rearrange('b t n d -> b t (n d)'),
+                Rearrange('... n d -> ... (n d)'),
                 Linear(num_latent_tokens * dim_latent, dim * num_spatial_tokens),
-                Rearrange('b t (s d) -> b t s d', s = num_spatial_tokens)
+                Rearrange('... (s d) -> ... s d', s = num_spatial_tokens)
             )
 
             self.to_latent_pred = Sequential(
                 RMSNorm(dim),
                 LinearNoBias(dim, dim_latent * latent_tokens_to_space),
-                Rearrange('b t s (n d) -> b t (s n) d', n = latent_tokens_to_space)
+                Rearrange('b t v s (n d) -> b t v (s n) d', n = latent_tokens_to_space)
+            )
+
+        # number of video views, for robotics, which could have third person + wrist camera at least
+
+        assert num_video_views >= 1
+        self.video_has_multi_view = num_video_views > 1
+
+        self.num_video_views = num_video_views
+
+        if self.video_has_multi_view:
+            self.view_emb = nn.Parameter(torch.randn(num_video_views, dim) * 1e-2)
+
+        # proprioception
+
+        self.has_proprio = exists(dim_proprio)
+        self.dim_proprio = dim_proprio
+
+        if self.has_proprio:
+            self.to_proprio_token = nn.Linear(dim_proprio, dim)
+
+            self.to_proprio_pred = Sequential(
+                RMSNorm(dim),
+                nn.Linear(dim, dim_proprio)
             )
 
         # register tokens
@@ -1597,6 +2162,7 @@ class DynamicsWorldModel(Module):
         # learned set of latent genes
 
         self.agent_has_genes = num_latent_genes > 0
+        self.num_latent_genes = num_latent_genes
         self.latent_genes = Parameter(randn(num_latent_genes, dim) * 1e-2)
 
         # policy head
@@ -1616,10 +2182,14 @@ class DynamicsWorldModel(Module):
             num_continuous_actions = num_continuous_actions,
             continuous_norm_stats = continuous_norm_stats,
             can_unembed = True,
-            unembed_dim = dim * 4
+            unembed_dim = dim * 4,
+            num_unembed_preds = multi_token_pred_len,
+            squeeze_unembed_preds = False
         )
 
-        self.behavior_clone_weight = behavior_clone_weight
+        # multi token prediction length
+
+        self.multi_token_pred_len = multi_token_pred_len
 
         # each agent token will have the reward embedding of the previous time step - but could eventually just give reward its own token
 
@@ -1632,12 +2202,15 @@ class DynamicsWorldModel(Module):
             learned_embedding = add_reward_embed_to_agent_token
         )
 
-        self.to_reward_pred = Sequential(
+        to_reward_pred = Sequential(
             RMSNorm(dim),
             LinearNoBias(dim, self.reward_encoder.num_bins)
         )
 
-        self.reward_loss_weight = reward_loss_weight
+        self.to_reward_pred = Ensemble(
+            to_reward_pred,
+            multi_token_pred_len
+        )
 
         # value head
 
@@ -1653,13 +2226,16 @@ class DynamicsWorldModel(Module):
         self.transformer = AxialSpaceTimeTransformer(
             dim = dim,
             depth = depth,
+            attn_heads = attn_heads,
             attn_dim_head = attn_dim_head,
             attn_softclamp_value = attn_softclamp_value,
             attn_kwargs = attn_kwargs,
             ff_kwargs = ff_kwargs,
             num_residual_streams = num_residual_streams,
             num_special_spatial_tokens = num_agents,
-            final_norm = False
+            time_block_every = time_block_every,
+            final_norm = False,
+            **transformer_kwargs
         )
 
         # ppo related
@@ -1670,9 +2246,40 @@ class DynamicsWorldModel(Module):
 
         self.ppo_eps_clip = ppo_eps_clip
         self.value_clip = value_clip
-        self.policy_entropy_weight = value_clip
+        self.policy_entropy_weight = policy_entropy_weight
+
+        # pmpo related
+
+        self.pmpo_pos_to_neg_weight = pmpo_pos_to_neg_weight
+        self.pmpo_kl_div_loss_weight = pmpo_kl_div_loss_weight
+        self.pmpo_reverse_kl = pmpo_reverse_kl
+
+        # rewards related
+
+        self.keep_reward_ema_stats = keep_reward_ema_stats
+        self.reward_ema_decay = reward_ema_decay
+
+        self.register_buffer('reward_quantile_filter', tensor(reward_quantile_filter), persistent = False)
+
+        self.register_buffer('ema_returns_mean', tensor(0.))
+        self.register_buffer('ema_returns_var', tensor(1.))
+
+        # loss related
+
+        self.flow_loss_normalizer = LossNormalizer(1)
+        self.reward_loss_normalizer = LossNormalizer(multi_token_pred_len)
+        self.discrete_actions_loss_normalizer = LossNormalizer(multi_token_pred_len) if num_discrete_actions > 0 else None
+        self.continuous_actions_loss_normalizer = LossNormalizer(multi_token_pred_len) if num_continuous_actions > 0 else None
+
+        self.latent_flow_loss_weight = latent_flow_loss_weight
+
+        self.register_buffer('reward_loss_weight', tensor(reward_loss_weight))
+        self.register_buffer('discrete_action_loss_weight', tensor(discrete_action_loss_weight))
+        self.register_buffer('continuous_action_loss_weight', tensor(continuous_action_loss_weight))
 
-        # zero
+        assert self.reward_loss_weight.numel() in {1, multi_token_pred_len}
+        assert self.discrete_action_loss_weight.numel() in {1, multi_token_pred_len}
+        assert self.continuous_action_loss_weight.numel() in {1, multi_token_pred_len}
 
         self.register_buffer('zero', tensor(0.), persistent = False)
 
@@ -1680,17 +2287,19 @@ class DynamicsWorldModel(Module):
     def device(self):
         return self.zero.device
 
-    def get_times_from_signal_level(
-        self,
-        signal_levels,
-        align_dims_left_to = None
-    ):
-        times = signal_levels.float() / self.max_steps
+    # types of parameters
 
-        if not exists(align_dims_left_to):
-            return times
+    def muon_parameters(self):
+        return self.transformer.muon_parameters()
 
-        return align_dims_left(times, align_dims_left_to)
+    def policy_head_parameters(self):
+        return [
+            *self.policy_head.parameters(),
+            *self.action_embedder.unembed_parameters() # includes the unembed from the action-embedder
+        ]
+
+    def value_head_parameters(self):
+        return self.value_head.parameters()
 
     def parameter(self):
         params = super().parameters()
@@ -1700,53 +2309,402 @@ class DynamicsWorldModel(Module):
 
         return list(set(params) - set(self.video_tokenizer.parameters()))
 
-    def learn_policy_from_generations(
+    # helpers for shortcut flow matching
+
+    def get_times_from_signal_level(
+        self,
+        signal_levels,
+        align_dims_left_to = None
+    ):
+        times = signal_levels.float() / self.max_steps
+
+        if not exists(align_dims_left_to):
+            return times
+
+        return align_dims_left(times, align_dims_left_to)
+
+    # evolutionary policy optimization - https://web3.arxiv.org/abs/2503.19037
+
+    @torch.no_grad()
+    def evolve_(
+        self,
+        fitness,
+        select_frac = 0.5,
+        tournament_frac = 0.5
+    ):
+        assert fitness.numel() == self.num_latent_genes
+
+        pop = self.latent_genes
+
+        pop_size = self.num_latent_genes
+        num_selected = ceil(pop_size * select_frac)
+        num_children = pop_size - num_selected
+
+        dim_gene = pop.shape[-1]
+
+        # natural selection just a sort and slice
+
+        selected_fitness, selected_indices = fitness.topk(num_selected, dim = -1)
+        selected = pop[selected_indices]
+
+        # use tournament - one tournament per child
+
+        tournament_size = max(2, ceil(num_selected * tournament_frac))
+
+        tournaments = torch.randn((num_children, num_selected), device = self.device).argsort(dim = -1)[:, :tournament_size]
+
+        parent_ids = selected_fitness[tournaments].topk(2, dim = -1).indices # get top 2 winners as parents
+
+        parents = selected[parent_ids]
+
+        # crossover by random interpolation from parent1 to parent2
+
+        random_uniform_mix = torch.randn((num_children, dim_gene), device = self.device).sigmoid()
+
+        parent1, parent2 = parents.unbind(dim = 1)
+        children = parent1.lerp(parent2, random_uniform_mix)
+
+        # store next population
+
+        next_pop = cat((selected, children))
+
+        self.latent_genes.copy_(next_pop)
+
+    # interacting with env for experience
+
+    @torch.no_grad()
+    def interact_with_env(
+        self,
+        env,
+        seed = None,
+        agent_index = 0,
+        step_size = 4,
+        max_timesteps = 16,
+        env_is_vectorized = False,
+        use_time_kv_cache = True,
+        store_agent_embed = True,
+        store_old_action_unembeds = True,
+    ):
+        assert exists(self.video_tokenizer)
+
+        init_frame = env.reset()
+
+        # frame to video
+
+        if env_is_vectorized:
+            video = rearrange(init_frame, 'b c vh vw -> b c 1 vh vw')
+        else:
+            video = rearrange(init_frame, 'c vh vw -> 1 c 1 vh vw')
+
+        batch, device = video.shape[0], video.device
+
+        # accumulate
+
+        rewards = None
+        discrete_actions = None
+        continuous_actions = None
+        discrete_log_probs = None
+        continuous_log_probs = None
+        values = None
+        latents = None
+
+        acc_agent_embed = None
+        acc_policy_embed = None
+
+        # keep track of termination, for setting the `is_truncated` field on Experience and for early stopping interaction with env
+
+        is_terminated = full((batch,), False, device = device)
+        is_truncated = full((batch,), False, device = device)
+
+        episode_lens = full((batch,), 0, device = device)
+
+        # maybe time kv cache
+
+        time_kv_cache = None
+
+        step_index = 0
+
+        while not is_terminated.all():
+            step_index += 1
+
+            latents = self.video_tokenizer(video, return_latents = True)
+
+            _, (agent_embed, next_time_kv_cache) = self.forward(
+                latents = latents,
+                signal_levels = self.max_steps - 1,
+                step_sizes = step_size,
+                rewards = rewards,
+                discrete_actions = discrete_actions,
+                continuous_actions = continuous_actions,
+                time_kv_cache = time_kv_cache,
+                latent_is_noised = True,
+                return_pred_only = True,
+                return_intermediates = True
+            )
+
+            # time kv cache
+
+            if use_time_kv_cache:
+                time_kv_cache = next_time_kv_cache
+
+            # get one agent
+
+            one_agent_embed = agent_embed[..., -1:, agent_index, :]
+
+            # values
+
+            value_bins = self.value_head(one_agent_embed)
+            value = self.reward_encoder.bins_to_scalar_value(value_bins)
+
+            values = safe_cat((values, value), dim = 1)
+
+            # policy embed
+
+            policy_embed = self.policy_head(one_agent_embed)
+
+            if store_old_action_unembeds:
+                acc_policy_embed = safe_cat((acc_policy_embed, policy_embed), dim = 1)
+
+            # sample actions
+
+            sampled_discrete_actions, sampled_continuous_actions = self.action_embedder.sample(policy_embed, pred_head_index = 0, squeeze = True)
+
+            discrete_actions = safe_cat((discrete_actions, sampled_discrete_actions), dim = 1)
+            continuous_actions = safe_cat((continuous_actions, sampled_continuous_actions), dim = 1)
+
+            # get the log prob and values for policy optimization
+
+            one_discrete_log_probs, one_continuous_log_probs = self.action_embedder.log_probs(
+                policy_embed,
+                pred_head_index = 0,
+                discrete_targets = sampled_discrete_actions,
+                continuous_targets = sampled_continuous_actions,
+            )
+
+            discrete_log_probs = safe_cat((discrete_log_probs, one_discrete_log_probs), dim = 1)
+            continuous_log_probs = safe_cat((continuous_log_probs, one_continuous_log_probs), dim = 1)
+
+            # pass the sampled action to the environment and get back next state and reward
+
+            env_step_out = env.step((sampled_discrete_actions, sampled_continuous_actions))
+
+            if len(env_step_out) == 2:
+                next_frame, reward = env_step_out
+                terminated = full((batch,), False)
+                truncated = full((batch,), False)
+
+            elif len(env_step_out) == 3:
+                next_frame, reward, terminated = env_step_out
+                truncated = full((batch,), False)
+
+            elif len(env_step_out) == 4:
+                next_frame, reward, terminated, truncated = env_step_out
+
+            elif len(env_step_out) == 5:
+                next_frame, reward, terminated, truncated, info = env_step_out
+
+            # update episode lens
+
+            episode_lens = torch.where(is_terminated, episode_lens, episode_lens + 1)
+
+            # update `is_terminated`
+
+            # (1) - environment says it is terminated
+            # (2) - previous step is truncated (this step is for bootstrap value)
+
+            is_terminated |= (terminated | is_truncated)
+
+            # update `is_truncated`
+
+            if step_index <= max_timesteps:
+                is_truncated |= truncated
+
+            if step_index == max_timesteps:
+                # if the step index is at the max time step allowed, set the truncated flag, if not already terminated
+
+                is_truncated |= ~is_terminated
+
+            # batch and time dimension
+
+            if env_is_vectorized:
+                next_frame = rearrange(next_frame, 'b c vh vw -> b c 1 vh vw')
+                reward = rearrange(reward, 'b -> b 1')
+            else:
+                next_frame = rearrange(next_frame, 'c vh vw -> 1 c 1 vh vw')
+                reward = rearrange(reward, ' -> 1 1')
+
+            # concat
+
+            video = cat((video, next_frame), dim = 2)
+            rewards = safe_cat((rewards, reward), dim = 1)
+
+            acc_agent_embed = safe_cat((acc_agent_embed, one_agent_embed), dim = 1)
+
+        # package up one experience for learning
+
+        batch, device = latents.shape[0], latents.device
+
+        one_experience = Experience(
+            latents = latents,
+            video = video[:, :, :-1],
+            rewards = rewards,
+            actions = (discrete_actions, continuous_actions),
+            log_probs = (discrete_log_probs, continuous_log_probs),
+            values = values,
+            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if exists(acc_policy_embed) and store_old_action_unembeds else None,
+            agent_embed = acc_agent_embed if store_agent_embed else None,
+            step_size = step_size,
+            agent_index = agent_index,
+            is_truncated = is_truncated,
+            lens = episode_lens,
+            is_from_world_model = False
+        )
+
+        return one_experience
+
+    # ppo
+
+    def learn_from_experience(
         self,
-        generation: Experience,
+        experience: Experience,
         policy_optim: Optimizer | None = None,
-        value_optim: Optimizer | None = None
+        value_optim: Optimizer | None = None,
+        only_learn_policy_value_heads = True, # in the paper, they do not finetune the entire dynamics model, they just learn the heads
+        use_pmpo = True,
+        normalize_advantages = None,
+        eps = 1e-6
     ):
-        latents = generation.latents
-        actions = generation.actions
-        old_log_probs = generation.log_probs
-        old_values = generation.values
-        rewards = generation.rewards
+        assert isinstance(experience, Experience)
+
+        experience = experience.to(self.device)
+
+        latents = experience.latents
+        actions = experience.actions
+        old_log_probs = experience.log_probs
+        old_values = experience.values
+        rewards = experience.rewards
+        agent_embeds = experience.agent_embed
+        old_action_unembeds = experience.old_action_unembeds
+
+        step_size = experience.step_size
+        agent_index = experience.agent_index
+
+        assert all([*map(exists, (old_log_probs, actions, old_values, rewards, step_size))]), 'the generations need to contain the log probs, values, and rewards for policy optimization - world_model.generate(..., return_log_probs_and_values = True)'
+
+        batch, time = latents.shape[0], latents.shape[1]
 
-        step_size = generation.step_size
-        agent_index = generation.agent_index
+        # calculate returns
 
-        assert all([*map(exists, (old_log_probs, actions, old_values, rewards, step_size))]), 'the generations need to contain the log probs, values, and rewards for policy optimization'
+        # mask out anything after the `lens`, which may include a bootstrapped node at the very end if `is_truncated = True`
+
+        if not exists(experience.is_truncated):
+            experience.is_truncated = full((batch,), True, device = latents.device)
+
+        if exists(experience.lens):
+            mask_for_gae = lens_to_mask(experience.lens, time)
+
+            rewards = rewards.masked_fill(~mask_for_gae, 0.)
+            old_values = old_values.masked_fill(~mask_for_gae, 0.)
+
+        # calculate returns
 
         returns = calc_gae(rewards, old_values, gamma = self.gae_discount_factor, lam = self.gae_lambda, use_accelerated = self.gae_use_accelerated)
 
-        # apparently they just use the sign of the advantage
+        # handle variable lengths
+
+        max_time = latents.shape[1]
+        is_var_len = exists(experience.lens)
+
+        mask = None
+
+        if is_var_len:
+            learnable_lens = experience.lens - experience.is_truncated.long() # if is truncated, remove the last one, as it is bootstrapped value
+            mask = lens_to_mask(learnable_lens, max_time)
+
+        # determine whether to finetune entire transformer or just learn the heads
+
+        world_model_forward_context = torch.no_grad if only_learn_policy_value_heads else nullcontext
+
+        # maybe keep track returns statistics and normalize returns and values before calculating advantage, as done in dreamer v3
+
+        if self.keep_reward_ema_stats:
+            ema_returns_mean, ema_returns_var = self.ema_returns_mean, self.ema_returns_var
+
+            decay = 1. - self.reward_ema_decay
+
+            # quantile filter
+
+            lo, hi = torch.quantile(returns, self.reward_quantile_filter).tolist()
+            returns_for_stats = returns.clamp(lo, hi)
+
+            # mean, var - todo - handle distributed
+
+            returns_mean, returns_var = returns_for_stats.mean(), returns_for_stats.var()
+
+            # ema
+
+            ema_returns_mean.lerp_(returns_mean, decay)
+            ema_returns_var.lerp_(returns_var, decay)
+
+            # normalize
+
+            ema_returns_std = ema_returns_var.clamp(min = 1e-5).sqrt()
+
+            normed_returns = (returns - ema_returns_mean) / ema_returns_std
+            normed_old_values = (old_values - ema_returns_mean) / ema_returns_std
+
+            advantage = normed_returns - normed_old_values
+        else:
+            advantage = returns - old_values
+
+        # if using pmpo, do not normalize advantages, but can be overridden
+
+        normalize_advantages = default(normalize_advantages, not use_pmpo)
+
+        if normalize_advantages:
+            advantage = F.layer_norm(advantage, advantage.shape, eps = eps)
+
         # https://arxiv.org/abs/2410.04166v1
 
-        advantage = (returns - old_values).sign()
+        if use_pmpo:
+            pos_advantage_mask = advantage >= 0.
+            neg_advantage_mask = ~pos_advantage_mask
 
         # replay for the action logits and values
+        # but only do so if fine tuning the entire world model for RL
 
         discrete_actions, continuous_actions = actions
 
-        _, agent_embed = self.forward(
-            latents = latents,
-            signal_levels = self.max_steps - 1,
-            step_sizes = step_size,
-            rewards = rewards,
-            discrete_actions = discrete_actions,
-            continuous_actions = continuous_actions,
-            latent_is_noised = True,
-            return_pred_only = True,
-            return_agent_tokens = True
-        )
+        if (
+            not only_learn_policy_value_heads or
+            not exists(agent_embeds)
+        ):
+
+            with world_model_forward_context():
+                _, (agent_embeds, _) = self.forward(
+                    latents = latents,
+                    signal_levels = self.max_steps - 1,
+                    step_sizes = step_size,
+                    rewards = rewards,
+                    discrete_actions = discrete_actions,
+                    continuous_actions = continuous_actions,
+                    latent_is_noised = True,
+                    return_pred_only = True,
+                    return_intermediates = True
+                )
+
+            agent_embeds = agent_embeds[..., agent_index, :]
 
-        agent_embed = agent_embed[..., agent_index, :]
+        # maybe detach agent embed
+
+        if only_learn_policy_value_heads:
+            agent_embeds = agent_embeds.detach()
 
         # ppo
 
-        policy_embed = self.policy_head(agent_embed)
+        policy_embed = self.policy_head(agent_embeds)
 
-        log_probs, entropies = self.action_embedder.log_probs(policy_embed, discrete_targets = discrete_actions, continuous_targets = continuous_actions, return_entropies = True)
+        log_probs, entropies = self.action_embedder.log_probs(policy_embed, pred_head_index = 0, discrete_targets = discrete_actions, continuous_targets = continuous_actions, return_entropies = True)
 
         # concat discrete and continuous actions into one for optimizing
 
@@ -1754,29 +2712,77 @@ class DynamicsWorldModel(Module):
         log_probs = safe_cat(log_probs, dim = -1)
         entropies = safe_cat(entropies, dim = -1)
 
-        ratio = (log_probs - old_log_probs).exp()
+        advantage = rearrange(advantage, '... -> ... 1') # broadcast across all actions
 
-        clipped_ratio = ratio.clamp(1. - self.ppo_eps_clip, 1. + self.ppo_eps_clip)
+        if use_pmpo:
+            # pmpo - weighting the positive and negative advantages equally - ignoring magnitude of advantage and taking the sign
+            # seems to be weighted across batch and time, iiuc
+            # eq (10) in https://arxiv.org/html/2410.04166v1
 
-        advantage = rearrange(advantage, '... -> ... 1') # broadcast across all actions
+            if exists(mask):
+                pos_advantage_mask &= mask
+                neg_advantage_mask &= mask
+
+            α = self.pmpo_pos_to_neg_weight
+
+            pos = masked_mean(log_probs, pos_advantage_mask)
+            neg = -masked_mean(log_probs, neg_advantage_mask)
+
+            policy_loss = -(α * pos + (1. - α) * neg)
+
+            # take care of kl
+
+            if self.pmpo_kl_div_loss_weight > 0.:
+
+                new_unembedded_actions = self.action_embedder.unembed(policy_embed, pred_head_index = 0)
+
+                kl_div_inputs, kl_div_targets = new_unembedded_actions, old_action_unembeds
+
+                # mentioned that the "reverse direction for the prior KL" was used
+                # make optional, as observed instability in toy task
+
+                if self.pmpo_reverse_kl:
+                    kl_div_inputs, kl_div_targets = kl_div_targets, kl_div_inputs
+
+                discrete_kl_div, continuous_kl_div = self.action_embedder.kl_div(kl_div_inputs, kl_div_targets)
 
-        # clipped surrogate loss
+                # accumulate discrete and continuous kl div
 
-        policy_loss = -torch.min(ratio * advantage, clipped_ratio * advantage)
-        policy_loss = reduce(policy_loss, 'b t na -> b t', 'sum')
+                kl_div_loss = 0.
 
-        policy_loss = policy_loss.mean()
+                if exists(discrete_kl_div):
+                    kl_div_loss = kl_div_loss + discrete_kl_div[mask].mean()
+
+                if exists(continuous_kl_div):
+                    kl_div_loss = kl_div_loss + continuous_kl_div[mask].mean()
+
+                policy_loss = policy_loss + kl_div_loss * self.pmpo_kl_div_loss_weight
+
+        else:
+            # ppo clipped surrogate loss
+
+            ratio = (log_probs - old_log_probs).exp()
+            clipped_ratio = ratio.clamp(1. - self.ppo_eps_clip, 1. + self.ppo_eps_clip)
+
+            policy_loss = -torch.min(ratio * advantage, clipped_ratio * advantage)
+            policy_loss = reduce(policy_loss, 'b t na -> b t', 'sum')
+
+            policy_loss = masked_mean(policy_loss, mask)
 
         # handle entropy loss for naive exploration bonus
 
-        entropy_loss = - reduce(entropies, 'b t na -> b t', 'sum').mean()
+        entropy_loss = - reduce(entropies, 'b t na -> b t', 'sum')
+
+        entropy_loss = masked_mean(entropy_loss, mask)
+
+        # total policy loss
 
         total_policy_loss = (
             policy_loss +
             entropy_loss * self.policy_entropy_weight
         )
 
-        # maye take policy optimizer step
+        # maybe take policy optimizer step
 
         if exists(policy_optim):
             total_policy_loss.backward()
@@ -1786,7 +2792,7 @@ class DynamicsWorldModel(Module):
 
         # value loss
 
-        value_bins = self.value_head(agent_embed)
+        value_bins = self.value_head(agent_embeds)
         values = self.reward_encoder.bins_to_scalar_value(value_bins)
 
         clipped_values = old_values + (values - old_values).clamp(-self.value_clip, self.value_clip)
@@ -1794,10 +2800,19 @@ class DynamicsWorldModel(Module):
 
         return_bins = self.reward_encoder(returns)
 
+        value_bins, return_bins, clipped_value_bins = tuple(rearrange(t, 'b t l -> b l t') for t in (value_bins, return_bins, clipped_value_bins))
+
         value_loss_1 = F.cross_entropy(value_bins, return_bins, reduction = 'none')
         value_loss_2 = F.cross_entropy(clipped_value_bins, return_bins, reduction = 'none')
 
-        value_loss = torch.maximum(value_loss_1, value_loss_2).mean()
+        value_loss = torch.maximum(value_loss_1, value_loss_2)
+
+        # maybe variable length
+
+        if is_var_len:
+            value_loss = value_loss[mask].mean()
+        else:
+            value_loss = value_loss.mean()
 
         # maybe take value optimizer step
 
@@ -1816,22 +2831,49 @@ class DynamicsWorldModel(Module):
         num_steps = 4,
         batch_size = 1,
         agent_index = 0,
+        tasks: int | Tensor | None = None,
+        latent_gene_ids = None,
         image_height = None,
         image_width = None,
         return_decoded_video = None,
         context_signal_noise = 0.1,       # they do a noising of the past, this was from an old diffusion world modeling paper from EPFL iirc
+        time_kv_cache: Tensor | None = None,
+        use_time_kv_cache = True,
         return_rewards_per_frame = False,
         return_agent_actions = False,
-        return_log_probs_and_values = False
+        return_log_probs_and_values = False,
+        return_for_policy_optimization = False,
+        return_time_kv_cache = False,
+        store_agent_embed = True,
+        store_old_action_unembeds = True
 
     ): # (b t n d) | (b c t h w)
 
+        # handy flag for returning generations for rl
+
+        if return_for_policy_optimization:
+            return_agent_actions |= True
+            return_log_probs_and_values |= True
+            return_rewards_per_frame |= True
+
+        # more variables
+
+        has_proprio = self.has_proprio
         was_training = self.training
         self.eval()
 
+        # validation
+
         assert log2(num_steps).is_integer(), f'number of steps {num_steps} must be a power of 2'
         assert 0 < num_steps <= self.max_steps, f'number of steps {num_steps} must be between 0 and {self.max_steps}'
 
+        if isinstance(tasks, int):
+            tasks = full((batch_size,), tasks, device = self.device)
+
+        assert not exists(tasks) or tasks.shape[0] == batch_size
+
+        # get state latent shape
+
         latent_shape = self.latent_shape
 
         # derive step size
@@ -1841,9 +2883,16 @@ class DynamicsWorldModel(Module):
         # denoising
         # teacher forcing to start with
 
-        latents = empty((batch_size, 0, *latent_shape), device = self.device)
+        latents = empty((batch_size, 0, self.num_video_views, *latent_shape), device = self.device)
+
+        past_latents_context_noise = latents.clone()
 
-        past_context_noise = latents.clone()
+        # maybe internal state
+
+        if has_proprio:
+            proprio = empty((batch_size, 0, self.dim_proprio), device = self.device)
+
+            past_proprio_context_noise = proprio.clone()
 
         # maybe return actions
 
@@ -1858,8 +2907,17 @@ class DynamicsWorldModel(Module):
         decoded_continuous_log_probs = None
         decoded_values = None
 
+        # maybe store agent embed
+
+        acc_agent_embed = None
+
+        # maybe store old actions for kl
+
+        acc_policy_embed = None
+
         # maybe return rewards
 
+        decoded_rewards = None
         if return_rewards_per_frame:
             decoded_rewards = empty((batch_size, 0), device = self.device, dtype = torch.float32)
 
@@ -1868,55 +2926,131 @@ class DynamicsWorldModel(Module):
         while latents.shape[1] < time_steps:
 
             curr_time_steps = latents.shape[1]
-            noised_latent = randn((batch_size, 1, *latent_shape), device = self.device)
 
-            for step in range(num_steps):
+            # determine whether to take an extra step if
+            # (1) using time kv cache
+            # (2) decoding anything off agent embedding (rewards, actions, etc)
+
+            take_extra_step = (
+                use_time_kv_cache or
+                return_rewards_per_frame or
+                store_agent_embed or
+                return_agent_actions
+            )
+
+            # prepare noised latent / proprio inputs
+
+            noised_latent = randn((batch_size, 1, self.num_video_views, *latent_shape), device = self.device)
+
+            noised_proprio = None
+
+            if has_proprio:
+                noised_proprio = randn((batch_size, 1, self.dim_proprio), device = self.device)
+
+            # denoising steps
+
+            for step in range(num_steps + int(take_extra_step)):
+
+                is_last_step = (step + 1) == num_steps
+
                 signal_levels = full((batch_size, 1), step * step_size, dtype = torch.long, device = self.device)
+ 
+                # noising past latent context
 
-                noised_context = latents.lerp(past_context_noise, context_signal_noise) # the paragraph after eq (8)
+                noised_context = latents.lerp(past_latents_context_noise, context_signal_noise) # the paragraph after eq (8)
 
-                noised_latent_with_context, pack_context_shape = pack((noised_context, noised_latent), 'b * n d')
+                noised_latent_with_context, pack_context_shape = pack((noised_context, noised_latent), 'b * v n d')
+
+                # handle proprio
+
+                noised_proprio_with_context = None
+
+                if has_proprio:
+                    noised_proprio_context = proprio.lerp(past_proprio_context_noise, context_signal_noise)
+                    noised_proprio_with_context, _ = pack((noised_proprio_context, noised_proprio), 'b * d')
+
+                # proper signal levels
 
                 signal_levels_with_context = F.pad(signal_levels, (curr_time_steps, 0), value = self.max_steps - 1)
 
-                pred, agent_embed = self.forward(
+                pred, (agent_embed, next_time_kv_cache) = self.forward(
                     latents = noised_latent_with_context,
                     signal_levels = signal_levels_with_context,
                     step_sizes = step_size,
                     rewards = decoded_rewards,
+                    tasks = tasks,
+                    latent_gene_ids = latent_gene_ids,
                     discrete_actions = decoded_discrete_actions,
                     continuous_actions = decoded_continuous_actions,
+                    proprio = noised_proprio_with_context,
+                    time_kv_cache = time_kv_cache,
                     latent_is_noised = True,
+                    latent_has_view_dim = True,
                     return_pred_only = True,
-                    return_agent_tokens = True
+                    return_intermediates = True,
                 )
 
-                _, pred = unpack(pred, pack_context_shape, 'b * n d')
+                if use_time_kv_cache and is_last_step:
+                    time_kv_cache = next_time_kv_cache
+
+                # early break if taking an extra step for agent embedding off cleaned latents for decoding
+
+                if take_extra_step and is_last_step:
+                    break
+
+                # maybe proprio
+
+                if has_proprio:
+                    pred, pred_proprio = pred
+
+                # unpack pred
+
+                _, pred = unpack(pred, pack_context_shape, 'b * v n d')
+
+                if has_proprio:
+                    _, pred_proprio = unpack(pred_proprio, pack_context_shape, 'b * d')
 
                 # derive flow, based on whether in x-space or not
 
-                if self.pred_orig_latent:
-                    times = self.get_times_from_signal_level(signal_levels, noised_latent)
-                    flow = (pred - noised_latent) / (1. - times)
-                else:
-                    flow = pred
+                def denoise_step(pred, noised, signal_levels):
+                    if self.pred_orig_latent:
+                        times = self.get_times_from_signal_level(signal_levels)
+                        aligned_times = align_dims_left(times, noised)
+
+                        flow = (pred - noised) / (1. - aligned_times)
+                    else:
+                        flow = pred
+
+                    return flow * (step_size / self.max_steps)
 
                 # denoise
 
-                noised_latent += flow * (step_size / self.max_steps)
+                noised_latent += denoise_step(pred, noised_latent, signal_levels)
+
+                if has_proprio:
+                    noised_proprio += denoise_step(pred_proprio, noised_proprio, signal_levels)
 
             denoised_latent = noised_latent # it is now denoised
 
+            if has_proprio:
+                denoised_proprio = noised_proprio
+
             # take care of the rewards by predicting on the agent token embedding on the last denoising step
 
             if return_rewards_per_frame:
                 one_agent_embed = agent_embed[:, -1:, agent_index]
 
-                reward_logits = self.to_reward_pred(one_agent_embed)
+                reward_logits = self.to_reward_pred.forward_one(one_agent_embed, id = 0)
                 pred_reward = self.reward_encoder.bins_to_scalar_value(reward_logits, normalize = True)
 
                 decoded_rewards = cat((decoded_rewards, pred_reward), dim = 1)
 
+            # maybe store agent embed
+
+            if store_agent_embed:
+                one_agent_embed = agent_embed[:, -1:, agent_index]
+                acc_agent_embed = safe_cat((acc_agent_embed, one_agent_embed), dim = 1)
+
             # decode the agent actions if needed
 
             if return_agent_actions:
@@ -1926,7 +3060,14 @@ class DynamicsWorldModel(Module):
 
                 policy_embed = self.policy_head(one_agent_embed)
 
-                sampled_discrete_actions, sampled_continuous_actions = self.action_embedder.sample(policy_embed)
+                # maybe store old actions
+
+                if store_old_action_unembeds:
+                    acc_policy_embed = safe_cat((acc_policy_embed, policy_embed), dim = 1)
+
+                # sample actions
+
+                sampled_discrete_actions, sampled_continuous_actions = self.action_embedder.sample(policy_embed, pred_head_index = 0, squeeze = True)
 
                 decoded_discrete_actions = safe_cat((decoded_discrete_actions, sampled_discrete_actions), dim = 1)
                 decoded_continuous_actions = safe_cat((decoded_continuous_actions, sampled_continuous_actions), dim = 1)
@@ -1934,6 +3075,7 @@ class DynamicsWorldModel(Module):
                 if return_log_probs_and_values:
                     discrete_log_probs, continuous_log_probs = self.action_embedder.log_probs(
                         policy_embed,
+                        pred_head_index = 0,
                         discrete_targets = sampled_discrete_actions,
                         continuous_targets = sampled_continuous_actions,
                     )
@@ -1952,7 +3094,14 @@ class DynamicsWorldModel(Module):
 
             # add new fixed context noise for the temporal consistency
 
-            past_context_noise = cat((past_context_noise, randn_like(denoised_latent)), dim = 1)
+            past_latents_context_noise = cat((past_latents_context_noise, randn_like(denoised_latent)), dim = 1)
+
+            # handle proprio
+
+            if has_proprio:
+                proprio = cat((proprio, denoised_proprio), dim = 1)
+
+                past_proprio_context_noise = cat((past_proprio_context_noise, randn_like(denoised_proprio)), dim = 1)
 
         # restore state
 
@@ -1966,24 +3115,50 @@ class DynamicsWorldModel(Module):
         video = None
 
         if return_decoded_video:
+
+            latents_for_video = rearrange(latents, 'b t v n d -> b v t n d')
+            latents_for_video, unpack_view = pack_one(latents_for_video, '* t n d')
+
             video = self.video_tokenizer.decode(
-                latents,
+                latents_for_video,
                 height = image_height,
                 width = image_width
             )
 
+            video = unpack_view(video, '* t c vh vw')
+
+        # remove the lone view dimension
+
+        if not self.video_has_multi_view:
+            latents = rearrange(latents, 'b t 1 ... -> b t ...')
+
+            if exists(video):
+                video = rearrange(video, 'b 1 ... -> b ...')
+
         # only return video or latent if not requesting anything else, for first stage training
 
-        if not has_at_least_one(return_rewards_per_frame, return_agent_actions):
-            return video if return_decoded_video else latents
+        if not has_at_least_one(return_rewards_per_frame, return_agent_actions, has_proprio):
+            out = video if return_decoded_video else latents
+
+            if not return_time_kv_cache:
+                return out
+
+            return out, time_kv_cache
 
         # returning agent actions, rewards, and log probs + values for policy optimization
 
+        batch, device = latents.shape[0], latents.device
+        experience_lens = full((batch,), time_steps, device = device)
+
         gen = Experience(
             latents = latents,
             video = video,
+            proprio = proprio if has_proprio else None,
+            agent_embed = acc_agent_embed if store_agent_embed else None,
+            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if exists(acc_policy_embed) and store_old_action_unembeds else None,
             step_size = step_size,
             agent_index = agent_index,
+            lens = experience_lens,
             is_from_world_model = True
         )
 
@@ -1998,13 +3173,17 @@ class DynamicsWorldModel(Module):
 
             gen.values = decoded_values
 
-        return gen
+        if not return_time_kv_cache:
+            return gen
+
+        return gen, time_kv_cache
 
     def forward(
         self,
         *,
-        video = None,                    # (b c t vh vw)
-        latents = None,                  # (b t n d) | (b t d)
+        video = None,                    # (b v? c t vh vw)
+        latents = None,                  # (b t v? n d) | (b t v? d)
+        lens = None,                     # (b)
         signal_levels = None,            # () | (b) | (b t)
         step_sizes = None,               # () | (b)
         step_sizes_log2 = None,          # () | (b)
@@ -2015,25 +3194,46 @@ class DynamicsWorldModel(Module):
         continuous_actions = None,       # (b t na) | (b t-1 na)
         discrete_action_types = None,    # (na)
         continuous_action_types = None,  # (na)
+        proprio = None,                  # (b t dp)
+        time_kv_cache = None,
         return_pred_only = False,
         latent_is_noised = False,
         return_all_losses = False,
-        return_agent_tokens = False,
-        add_autoregressive_action_loss = False
+        return_intermediates = False,
+        add_autoregressive_action_loss = True,
+        update_loss_ema = None,
+        latent_has_view_dim = False
     ):
         # handle video or latents
 
         assert exists(video) ^ exists(latents)
 
+        # standardize view dimension
+
+        if not self.video_has_multi_view:
+            if exists(video):
+                video = rearrange(video, 'b ... -> b 1 ...')
+
+            if exists(latents) and not latent_has_view_dim:
+                latents = rearrange(latents, 'b t ... -> b t 1 ...')
+
+        # if raw video passed in, tokenize
+
         if exists(video):
+            assert video.ndim == 6
+
+            video, unpack_views = pack_one(video, '* c t vh vw')
             assert exists(self.video_tokenizer), 'video_tokenizer must be passed in if training from raw video on dynamics model'
 
             latents = self.video_tokenizer.tokenize(video)
+            latents = unpack_views(latents, '* t n d')
+            latents = rearrange(latents, 'b v t n d -> b t v n d')
 
-        if latents.ndim == 3:
-            latents = rearrange(latents, 'b t d -> b t 1 d') # 1 latent edge case
+        if latents.ndim == 4:
+            latents = rearrange(latents, 'b t v d -> b t v 1 d') # 1 latent edge case
 
-        assert latents.shape[-2:] == self.latent_shape
+        assert latents.shape[-2:] == self.latent_shape, f'latents must have shape {self.latent_shape}, got {latents.shape[-2:]}'
+        assert latents.shape[2] == self.num_video_views, f'latents must have {self.num_video_views} views, got {latents.shape[2]}'
 
         # variables
 
@@ -2108,16 +3308,17 @@ class DynamicsWorldModel(Module):
 
         # times is from 0 to 1
 
-        times = self.get_times_from_signal_level(signal_levels, latents)
+        times = self.get_times_from_signal_level(signal_levels)
 
         if not latent_is_noised:
             # get the noise
 
             noise = randn_like(latents)
+            aligned_times = align_dims_left(times, latents)
 
             # noise from 0 as noise to 1 as data
 
-            noised_latents = noise.lerp(latents, times)
+            noised_latents = noise.lerp(latents, aligned_times)
 
         else:
             noised_latents = latents
@@ -2165,6 +3366,27 @@ class DynamicsWorldModel(Module):
 
                 reward_tokens = add('1 d, b t d', self.reward_learned_embed, reward_tokens)
 
+        # maybe proprioception
+
+        assert xnor(self.has_proprio, exists(proprio)), 'proprio must be passed in if `dim_proprio` is set and vice versa'
+
+        noised_proprio = None
+
+        if self.has_proprio:
+
+            if not latent_is_noised:
+                # get the noise
+
+                proprio_noise = randn_like(proprio)
+                aligned_times = align_dims_left(times, proprio)
+
+                # noise from 0 as noise to 1 as data
+
+                noised_proprio = proprio_noise.lerp(proprio, aligned_times)
+
+            else:
+                noised_proprio = proprio
+
         # maybe create the action tokens
 
         if exists(discrete_actions) or exists(continuous_actions):
@@ -2185,16 +3407,27 @@ class DynamicsWorldModel(Module):
 
             action_tokens = add('1 d, b t d', self.action_learned_embed, action_tokens)
 
+        elif self.action_embedder.has_actions:
+            action_tokens = torch.zeros_like(agent_tokens[:, :, 0:1])
+
         else:
             action_tokens = agent_tokens[:, :, 0:0] # else empty off agent tokens
 
         # main function, needs to be defined as such for shortcut training - additional calls for consistency loss
 
-        def get_prediction(noised_latents, signal_levels, step_sizes_log2, action_tokens, reward_tokens, agent_tokens, return_agent_tokens = False):
+        def get_prediction(noised_latents, noised_proprio, signal_levels, step_sizes_log2, action_tokens, reward_tokens, agent_tokens, return_agent_tokens = False, return_time_kv_cache = False):
+
             # latents to spatial tokens
 
             space_tokens = self.latents_to_spatial_tokens(noised_latents)
 
+            # maybe add view embedding
+
+            if self.video_has_multi_view:
+                space_tokens = add('b t v ... d, v d', space_tokens, self.view_emb)
+
+            # merge spatial tokens
+
             space_tokens, inverse_pack_space_per_latent = pack_one(space_tokens, 'b t * d')
 
             num_spatial_tokens = space_tokens.shape[-2]
@@ -2212,6 +3445,13 @@ class DynamicsWorldModel(Module):
 
             registers = repeat(self.register_tokens, 's d -> b t s d', b = batch, t = time)
 
+            # maybe proprio
+
+            if exists(noised_proprio):
+                proprio_token = self.to_proprio_token(noised_proprio)
+            else:
+                proprio_token = registers[:, :, 0:0]
+
             # determine signal + step size embed for their diffusion forcing + shortcut
 
             signal_embed = self.signal_levels_embed(signal_levels)
@@ -2224,15 +3464,15 @@ class DynamicsWorldModel(Module):
 
             # pack to tokens for attending
 
-            tokens, packed_tokens_shape = pack([flow_token, space_tokens, registers, action_tokens, reward_tokens, agent_tokens], 'b t * d')
+            tokens, packed_tokens_shape = pack([flow_token, space_tokens, proprio_token, registers, action_tokens, reward_tokens, agent_tokens], 'b t * d')
 
             # attention
 
-            tokens = self.transformer(tokens)
+            tokens, (next_time_kv_cache, *_) = self.transformer(tokens, kv_cache = time_kv_cache, return_intermediates = True)
 
             # unpack
 
-            flow_token, space_tokens, register_tokens, action_tokens, reward_tokens, agent_tokens = unpack(tokens, packed_tokens_shape, 'b t * d')
+            flow_token, space_tokens, proprio_token, register_tokens, action_tokens, reward_tokens, agent_tokens = unpack(tokens, packed_tokens_shape, 'b t * d')
 
             # pooling
 
@@ -2240,10 +3480,22 @@ class DynamicsWorldModel(Module):
 
             pred = self.to_latent_pred(space_tokens)
 
+            # maybe proprio
+
+            if self.has_proprio:
+                pred_proprio = self.to_proprio_pred(proprio_token)
+
+                pred = (pred, pred_proprio)
+
+            # returning
+
             if not return_agent_tokens:
                 return pred
 
-            return pred, agent_tokens
+            if not return_time_kv_cache:
+                return pred, agent_tokens
+
+            return pred, (agent_tokens, next_time_kv_cache)
 
         # curry into get_prediction what does not change during first call as well as the shortcut ones
 
@@ -2251,13 +3503,47 @@ class DynamicsWorldModel(Module):
 
         # forward the network
 
-        pred, encoded_agent_tokens = _get_prediction(noised_latents, signal_levels, step_sizes_log2, return_agent_tokens = True)
+        pred, (encoded_agent_tokens, next_time_kv_cache) = _get_prediction(noised_latents, noised_proprio, signal_levels, step_sizes_log2, return_agent_tokens = True, return_time_kv_cache = True)
 
         if return_pred_only:
-            if not return_agent_tokens:
+            if not return_intermediates:
+                return pred
+
+            return pred, (encoded_agent_tokens, next_time_kv_cache)
+
+        # pack the predictions to calculate flow for different modalities all at once
+
+        if self.has_proprio:
+            pred, for_flow_loss_packed_shape = pack(pred, 'b t *')
+
+            noised, _ = pack((noised_latents, noised_proprio), 'b t *')
+            data, _ = pack((latents, proprio), 'b t *')
+            noise, _ = pack((noise, proprio_noise), 'b t *')
+        else:
+            noised = noised_latents
+            data = latents
+
+        # wrapper function for maybe unpacking and packing modalities for doing flow math in unison
+
+        def maybe_pack_unpack(fn):
+            @wraps(fn)
+            @torch.no_grad()
+            def inner(noised, *args, **kwargs):
+
+                noised_proprio = None
+
+                if self.has_proprio:
+                    noised, noised_proprio = unpack(noised, for_flow_loss_packed_shape, 'b t *')
+
+                pred = fn(noised, noised_proprio, *args, **kwargs)
+
+                if self.has_proprio:
+                    pred, _ = pack(pred, 'b t *')
+
                 return pred
+            return inner
 
-            return pred, encoded_agent_tokens
+        wrapped_get_prediction = maybe_pack_unpack(_get_prediction)
 
         # determine the target for the loss
 
@@ -2274,46 +3560,45 @@ class DynamicsWorldModel(Module):
             # x-space as in paper is in else clause
 
             if is_v_space_pred:
-                pred_target = flow = latents - noise
+                pred_target = flow = data - noise
             else:
-                pred_target = latents
+                pred_target = data
         else:
             # shortcut training - Frans et al. https://arxiv.org/abs/2410.12557
 
             # basically a consistency loss where you ensure quantity of two half steps equals one step
             # dreamer then makes it works for x-space with some math
 
-            get_prediction_no_grad = torch.no_grad()(_get_prediction)
-
             step_sizes_log2_minus_one = step_sizes_log2 - 1 # which equals d / 2
             half_step_size = 2 ** step_sizes_log2_minus_one
 
-            first_step_pred = get_prediction_no_grad(noised_latents, signal_levels, step_sizes_log2_minus_one)
+            first_step_pred = wrapped_get_prediction(noised, signal_levels, step_sizes_log2_minus_one)
 
             # first derive b'
 
             if is_v_space_pred:
                 first_step_pred_flow = first_step_pred
             else:
-                first_times = self.get_times_from_signal_level(signal_levels, noised_latents)
-                first_step_pred_flow = (first_step_pred - noised_latents) / (1. - first_times)
+                first_times = self.get_times_from_signal_level(signal_levels, noised)
+
+                first_step_pred_flow = (first_step_pred - noised) / (1. - first_times)
 
             # take a half step
 
-            half_step_size_align_left = align_dims_left(half_step_size, noised_latents)
+            half_step_size_align_left = align_dims_left(half_step_size, noised)
 
-            denoised_latent = noised_latents + first_step_pred_flow * (half_step_size_align_left / self.max_steps)
+            denoised = noised + first_step_pred_flow * (half_step_size_align_left / self.max_steps)
 
             # get second prediction for b''
 
             signal_levels_plus_half_step = signal_levels + half_step_size[:, None]
-            second_step_pred = get_prediction_no_grad(denoised_latent, signal_levels_plus_half_step, step_sizes_log2_minus_one)
+            second_step_pred = wrapped_get_prediction(denoised, signal_levels_plus_half_step, step_sizes_log2_minus_one)
 
             if is_v_space_pred:
                 second_step_pred_flow = second_step_pred
             else:
-                second_times = self.get_times_from_signal_level(signal_levels_plus_half_step, denoised_latent)
-                second_step_pred_flow = (second_step_pred - denoised_latent) / (1. - second_times)
+                second_times = self.get_times_from_signal_level(signal_levels_plus_half_step, denoised)
+                second_step_pred_flow = (second_step_pred - denoised) / (1. - second_times)
 
             # pred target is sg(b' + b'') / 2
 
@@ -2322,7 +3607,7 @@ class DynamicsWorldModel(Module):
             # need to convert x-space to v-space
 
             if is_x_space:
-                pred = (pred - noised_latents) / (1. - first_times)
+                pred = (pred - noised) / (1. - first_times)
                 maybe_shortcut_loss_weight = (1. - first_times) ** 2
 
         # mse loss
@@ -2335,9 +3620,23 @@ class DynamicsWorldModel(Module):
 
         if exists(self.loss_weight_fn):
             loss_weight = self.loss_weight_fn(times)
+            loss_weight = align_dims_left(loss_weight, flow_losses)
+
             flow_losses = flow_losses * loss_weight
 
-        flow_loss = flow_losses.mean()
+        # handle variable lengths if needed
+
+        is_var_len = exists(lens)
+
+        if is_var_len:
+
+            loss_mask = lens_to_mask(lens, time)
+            loss_mask_without_last = loss_mask[:, :-1]
+
+            flow_loss = flow_losses[loss_mask].mean()
+
+        else:
+            flow_loss = flow_losses.mean()
 
         # now take care of the agent token losses
 
@@ -2348,58 +3647,114 @@ class DynamicsWorldModel(Module):
             if rewards.ndim == 2: # (b t)
                 encoded_agent_tokens = reduce(encoded_agent_tokens, 'b t g d -> b t d', 'mean')
 
-            reward_pred = self.to_reward_pred(encoded_agent_tokens)
-            reward_loss = F.cross_entropy(reward_pred, two_hot_encoding)
+            reward_pred = self.to_reward_pred(encoded_agent_tokens[:, :-1])
+
+            reward_pred = rearrange(reward_pred, 'mtp b t l -> b l t mtp')
+
+            reward_targets, reward_loss_mask = create_multi_token_prediction_targets(two_hot_encoding[:, :-1], self.multi_token_pred_len)
+
+            reward_targets = rearrange(reward_targets, 'b t mtp l -> b l t mtp')
+
+            reward_losses = F.cross_entropy(reward_pred, reward_targets, reduction = 'none')
+
+            reward_losses = reward_losses.masked_fill(~reward_loss_mask, 0.)
+
+            if is_var_len:
+                reward_loss = reward_losses[loss_mask_without_last].mean(dim = 0)
+            else:
+                reward_loss = reduce(reward_losses, '... mtp -> mtp', 'mean') # they sum across the prediction steps (mtp dimension) - eq(9)
 
         # maybe autoregressive action loss
 
-        behavior_clone_loss = self.zero
+        discrete_action_loss = self.zero
+        continuous_action_loss = self.zero
 
         if (
             self.num_agents == 1 and
             add_autoregressive_action_loss and
+            time > 1,
             (exists(discrete_actions) or exists(continuous_actions))
         ):
             assert self.action_embedder.has_actions
 
+            # handle actions having time vs time - 1 length
+            # remove the first action if it is equal to time (as it would come from some agent token in the past)
+
+            if exists(discrete_actions) and discrete_actions.shape[1] == time:
+                discrete_actions = discrete_actions[:, 1:]
+
+            if exists(continuous_actions) and continuous_actions.shape[1] == time:
+                continuous_actions = continuous_actions[:, 1:]
+
             # only for 1 agent
 
             agent_tokens = rearrange(agent_tokens, 'b t 1 d -> b t d')
             policy_embed = self.policy_head(agent_tokens[:, :-1])
 
+            # constitute multi token prediction targets
+
+            discrete_action_targets = continuous_action_targets = None
+
+            if exists(discrete_actions):
+                discrete_action_targets, discrete_mask = create_multi_token_prediction_targets(discrete_actions, self.multi_token_pred_len)
+                discrete_action_targets = rearrange(discrete_action_targets, 'b t mtp ... -> mtp b t ...')
+                discrete_mask = rearrange(discrete_mask, 'b t mtp -> mtp b t')
+
+            if exists(continuous_actions):
+                continuous_action_targets, continuous_mask = create_multi_token_prediction_targets(continuous_actions, self.multi_token_pred_len)
+                continuous_action_targets = rearrange(continuous_action_targets, 'b t mtp ... -> mtp b t ...')
+                continuous_mask = rearrange(continuous_mask, 'b t mtp -> mtp b t')
+
             discrete_log_probs, continuous_log_probs = self.action_embedder.log_probs(
                 policy_embed,
-                discrete_targets = discrete_actions[:, 1:] if exists(discrete_actions) else None,
-                continuous_targets = continuous_actions[:, 1:] if exists(continuous_actions) else None
+                discrete_targets = discrete_action_targets if exists(discrete_actions) else None,
+                continuous_targets = continuous_action_targets if exists(continuous_actions) else None
             )
 
             if exists(discrete_log_probs):
-                behavior_clone_loss = behavior_clone_loss + discrete_log_probs.sum(dim = -1).mean()
+                discrete_log_probs = discrete_log_probs.masked_fill(~discrete_mask[..., None], 0.)
+
+                if is_var_len:
+                    discrete_action_losses = rearrange(-discrete_log_probs, 'mtp b t na -> b t na mtp')
+                    discrete_action_loss = reduce(discrete_action_losses[loss_mask_without_last], '... mtp -> mtp', 'mean')
+                else:
+                    discrete_action_loss = reduce(-discrete_log_probs, 'mtp b t na -> mtp', 'mean')
 
             if exists(continuous_log_probs):
-                behavior_clone_loss = behavior_clone_loss + continuous_log_probs.sum(dim = -1).mean()
+                continuous_log_probs = continuous_log_probs.masked_fill(~continuous_mask[..., None], 0.)
+
+                if is_var_len:
+                    continuous_action_losses = rearrange(-continuous_log_probs, 'mtp b t na -> b t na mtp')
+                    continuous_action_loss = reduce(continuous_action_losses[loss_mask_without_last], '... mtp -> mtp', 'mean')
+                else:
+                    continuous_action_loss = reduce(-continuous_log_probs, 'mtp b t na -> mtp', 'mean')
+
+        # handle loss normalization
+
+        losses = WorldModelLosses(flow_loss, reward_loss, discrete_action_loss, continuous_action_loss)
+
+        if exists(self.flow_loss_normalizer):
+            flow_loss = self.flow_loss_normalizer(flow_loss, update_ema = update_loss_ema)
 
-        # gather losses
+        if exists(rewards) and exists(self.reward_loss_normalizer):
+            reward_loss = self.reward_loss_normalizer(reward_loss, update_ema = update_loss_ema)
+
+        if exists(discrete_actions) and exists(self.discrete_actions_loss_normalizer):
+            discrete_action_loss = self.discrete_actions_loss_normalizer(discrete_action_loss, update_ema = update_loss_ema)
+
+        if exists(continuous_actions) and exists(self.continuous_actions_loss_normalizer):
+            continuous_action_loss = self.continuous_actions_loss_normalizer(continuous_action_loss, update_ema = update_loss_ema)
+
+        # gather losses - they sum across the multi token prediction steps for rewards and actions - eq (9)
 
         total_loss = (
-            flow_loss +
-            reward_loss * self.reward_loss_weight +
-            behavior_clone_loss * self.behavior_clone_weight
+            flow_loss * self.latent_flow_loss_weight +
+            (reward_loss * self.reward_loss_weight).sum() +
+            (discrete_action_loss * self.discrete_action_loss_weight).sum() + 
+            (continuous_action_loss * self.continuous_action_loss_weight).sum()
         )
 
         if not return_all_losses:
             return total_loss
 
-        losses = WorldModelLosses(flow_loss, reward_loss, behavior_clone_loss)
-
         return total_loss, losses
-
-# dreamer
-
-class Dreamer(Module):
-    def __init__(
-        self,
-        video_tokenizer: VideoTokenizer,
-        dynamics_model: DynamicsWorldModel,
-    ):
-        super().__init__()