Add files using upload-large-folder tool

bridge/models/ordinary200/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

bridge/models/ordinary200/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

bridge/preprocessing/action_ranges.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4492b53758cea8b3d4cb01e01ddf32d9d4bf31095fb3ca768a2d3ff631be8b15
+size 829

code/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,9 @@
+# Third-party provenance
+
+This is a local research-code consolidation, not a public release.
+
+- `objective.py` preserves the original ByteDance/Hugging Face Apache-2.0 header from the audited RLVR-World/verl-derived implementation.
+- `reward_core.py` preserves its original upstream header. The audited upstream RLVR-World reference is commit `e1b8b6f40ca0696919ce9b8dbe965c4153bcf5f1`.
+- Tokenizer and backbone loading in the archived trainers rely on iVideoGPT/RLVR-World components, Hugging Face Transformers and PyTorch. Those model assets and datasets are not included here.
+- The local source manifest records file-level provenance. Before public distribution, consolidate the applicable full upstream licenses and choose the license for original additions; no blanket relicensing is implied here.
+

code/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "futureworlds"
+version = "0.1.0.dev0"
+description = "Memory-conditioned search and post-training for robotic world models"
+requires-python = ">=3.9"
+dependencies = ["numpy", "torch>=2.1", "transformers>=4.45,<5", "safetensors"]
+
+[project.optional-dependencies]
+hub = ["huggingface_hub>=0.34"]
+
+[tool.setuptools.packages.find]
+where = ["src"]

code/src/futureworlds/__init__.py

@@ -0,0 +1,2 @@
+"""FutureWorlds research core. Training entry points are being consolidated."""
+__version__ = "0.1.0.dev0"

code/src/futureworlds/batch.py

@@ -0,0 +1,9 @@
+"""Small tensor container for the upstream video reward's batch-only API."""
+class TensorBatch:
+    def __init__(self,batch,meta_info=None):
+        self.batch=batch
+        self.meta_info=meta_info or {}
+    @classmethod
+    def from_single_dict(cls,batch,meta_info=None):return cls(batch,meta_info)
+    def __len__(self):return len(next(iter(self.batch.values())))
+    def __getitem__(self,index):return type(self)({k:v[index] for k,v in self.batch.items()},self.meta_info)

code/src/futureworlds/checkpoints.py

@@ -0,0 +1,76 @@
+"""Load an explicit released world-model variant; never substitute adapters."""
+import hashlib
+import json
+from pathlib import Path, PurePosixPath
+
+
+def safe_path(root, relative):
+    rel = PurePosixPath(relative)
+    if rel.is_absolute() or not rel.parts or '..' in rel.parts or '\\' in relative:
+        raise ValueError(f'Invalid relative path: {relative!r}')
+    root = Path(root).resolve()
+    path = (root / relative).resolve()
+    if path != root and root not in path.parents:
+        raise ValueError('Path escapes bundle root')
+    return path
+
+
+def file_sha256(path):
+    value = hashlib.sha256()
+    with Path(path).open('rb') as stream:
+        for block in iter(lambda: stream.read(8 * 1024 * 1024), b''):
+            value.update(block)
+    return value.hexdigest()
+
+
+def verify_bundle(root):
+    root = Path(root)
+    manifest = json.loads((root / 'manifest.json').read_text())
+    if manifest.get('schema_version') != 1 or not manifest.get('files'):
+        raise ValueError('Unsupported or empty manifest')
+    seen = set()
+    for item in manifest['files']:
+        name = item['path']
+        if name in seen:
+            raise ValueError(f'Duplicate file: {name}')
+        seen.add(name)
+        path = safe_path(root, name)
+        if not path.is_file() or path.stat().st_size != item['bytes']:
+            raise ValueError(f'Missing file or wrong size: {name}')
+        if file_sha256(path) != item['sha256']:
+            raise ValueError(f'Checksum mismatch: {name}')
+    return manifest
+
+
+def load_world_model(root, dataset, variant='main', device='cpu', verify=True):
+    """Load the predictor only. Codec/actions/T5 preprocessing are separate.
+
+    FP32 and SDPA match the audited loaders. Dataset/variant selection is
+    explicit; missing text adapters are an error, never silently zero-filled.
+    """
+    import torch
+    from safetensors.torch import load_file
+    from transformers import AutoConfig, AutoModelForCausalLM
+    from .text_conditioning import TextWorldModel
+
+    root = Path(root)
+    if verify:
+        verify_bundle(root)
+    settings = json.loads(safe_path(root, f'{dataset}/futureworlds_config.json').read_text())
+    if settings['dataset'] != dataset or variant not in settings['variants']:
+        raise ValueError('Dataset or variant does not match bundle configuration')
+    entry = settings['variants'][variant]
+    folder = safe_path(root, entry['path'])
+    cfg = AutoConfig.from_pretrained(folder, local_files_only=True, trust_remote_code=False)
+    core = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.float32, attn_implementation='sdpa', trust_remote_code=False)
+    core.load_state_dict(load_file(str(folder / 'model.safetensors')), strict=True)
+    if settings['text_conditioning']['enabled']:
+        tc = settings['text_conditioning']
+        model = TextWorldModel(core, text_width=tc['width'], layers=tc['layers'], inner=tc['inner'])
+        adapters = folder / 'text_adapters.safetensors'
+        if not adapters.is_file():
+            raise FileNotFoundError(f'Required trained text adapters missing: {adapters}')
+        model.adapters.load_state_dict(load_file(str(adapters)), strict=True)
+    else:
+        model = core
+    return model.to(device).eval()

code/src/futureworlds/decoding.py

@@ -0,0 +1,81 @@
+"""Fixed-length, constrained HF beam search and independent numerical checks."""
+import torch
+from transformers import GenerationConfig, LogitsProcessor, LogitsProcessorList
+
+class VisualTokensOnly(LogitsProcessor):
+    def __init__(self, valid_size):
+        self.valid_size = valid_size
+
+    def __call__(self, input_ids, scores):
+        scores = scores.clone()
+        scores[:, self.valid_size:] = -torch.inf
+        return scores
+
+@torch.inference_mode()
+def beam_decode(model, prefix, width=4, steps=80, valid_size=4375):
+    config = GenerationConfig(
+        do_sample=False, num_beams=width, num_return_sequences=width,
+        max_new_tokens=steps, min_new_tokens=steps,
+        eos_token_id=None, bos_token_id=None, pad_token_id=0,
+        length_penalty=0.0 if width > 1 else 1.0, early_stopping=False, renormalize_logits=True,
+        repetition_penalty=1.0, no_repeat_ngram_size=0,
+        use_cache=True, return_dict_in_generate=True, output_scores=True,
+    )
+    result = model.generate(
+        input_ids=prefix, attention_mask=torch.ones_like(prefix),
+        generation_config=config,
+        logits_processor=LogitsProcessorList([VisualTokensOnly(valid_size)]),
+        eos_token_id=None, bos_token_id=None,
+        forced_bos_token_id=None, forced_eos_token_id=None,
+    )
+    candidates = result.sequences[:, prefix.shape[1]:]
+    assert candidates.shape == (width,steps)
+    assert int(candidates.min()) >= 0 and int(candidates.max()) < valid_size
+    assert len({tuple(x) for x in candidates.tolist()}) == width
+    transitions = model.compute_transition_scores(
+        result.sequences, result.scores, getattr(result,'beam_indices',None),
+        normalize_logits=False,
+    )
+    assert transitions.shape == (width,steps) and torch.isfinite(transitions).all()
+    sums = transitions.sum(-1)
+    scores = result.sequences_scores if width > 1 else sums
+    error = float((sums-scores).abs().max())
+    assert torch.allclose(sums,scores,rtol=0,atol=2e-4), error
+    assert torch.all(scores[:-1] >= scores[1:])
+    return candidates, scores, {'beam_width':width,'tokens_per_frame':steps,
+        'selected_index':0,'candidate_scores':scores.tolist(),
+        'score_reconstruction_max_error':error,
+        'score_definition':'Sum of log probabilities renormalized over valid visual tokens only; no length penalty or GT/image reward.'}
+
+@torch.inference_mode()
+def independent_check():
+    from transformers import LlamaConfig, LlamaForCausalLM
+    torch.manual_seed(3719)
+    model = LlamaForCausalLM(LlamaConfig(vocab_size=8,hidden_size=32,
+        intermediate_size=48,num_hidden_layers=2,num_attention_heads=4,
+        num_key_value_heads=4,max_position_embeddings=128,
+        eos_token_id=2,bos_token_id=1,pad_token_id=0,
+        attention_dropout=0.0)).float().eval()
+    prefix = torch.tensor([[6,1,3]],dtype=torch.long)
+    reports=[]
+    for width in [1,4]:
+        beams=[([],0.0)]
+        for _ in range(4):
+            expanded=[]
+            for seq,score in beams:
+                x=torch.tensor([prefix[0].tolist()+seq],dtype=torch.long)
+                probs=model(input_ids=x,use_cache=False).logits[0,-1,:5].log_softmax(-1)
+                expanded.extend((seq+[token],score+float(probs[token])) for token in range(5))
+            beams=sorted(expanded,key=lambda item:item[1],reverse=True)[:width]
+        ids,scores,meta=beam_decode(model,prefix,width=width,steps=4,valid_size=5)
+        expected=torch.tensor([seq for seq,_ in beams],dtype=torch.long)
+        expected_scores=torch.tensor([score for _,score in beams])
+        assert torch.equal(ids,expected),(width,ids,expected)
+        assert torch.allclose(scores,expected_scores,rtol=0,atol=2e-5)
+        reports.append({'width':width,'matches_independent_uncached_search':True,
+            'max_score_error':float((scores-expected_scores).abs().max())})
+    return {'tests':reports,'eos_disabled':True,'valid_token_constraint_checked':True,'complete':True}
+
+if __name__ == '__main__':
+    import json, transformers
+    print(json.dumps({'transformers':transformers.__version__,**independent_check()}))

code/src/futureworlds/memory.py

@@ -0,0 +1,41 @@
+"""Causal ReCAP prompts shared by rollout and every policy log-probability pass."""
+import torch
+
+def retained(target,recent=6):
+    assert target>=2 and recent>=1
+    return [1]+list(range(max(2,target-recent),target))
+
+def make_prefix(context,history,actions,target,recent=6):
+    kept=retained(target,recent)
+    assert max(kept)==target-1 and len(kept)==len(set(kept))
+    blocks=[context]
+    for t in kept:blocks += [history[t],actions[t].expand(context.shape[0],-1)]
+    return torch.cat(blocks,dim=1),kept
+
+@torch.no_grad()
+def rollout(model,context,anchor,actions,horizon,group,generator,visual=4375,width=80,temperature=1.0):
+    model.eval()
+    context=context.expand(group,-1);history={1:anchor.expand(group,-1)}
+    frames=[];traces=[]
+    for target in range(2,horizon+2):
+        prefix,kept=make_prefix(context,history,actions,target)
+        pending=prefix;kv=None;codes=[];lps=[]
+        for token in range(width):
+            out=model(input_ids=pending,past_key_values=kv,use_cache=True)
+            kv=out.past_key_values
+            logp=torch.log_softmax(out.logits[:,-1,:visual].float()/temperature,-1)
+            pending=torch.multinomial(logp.exp(),1,generator=generator)
+            codes.append(pending);lps.append(logp.gather(1,pending).squeeze(1))
+        response=torch.cat(codes,1)
+        frames.append({'prefix':prefix.detach(),'response':response.detach(),'sample_logp':torch.stack(lps,1).detach(),'target':target})
+        history[target]=response
+        traces.append({'target':target,'retained':kept,'prefix_length':prefix.shape[1],'positions_reset':True,'kv_reset':True,'future_observations':False})
+    return frames,traces
+
+def frame_logp(model,prefix,response,visual=4375,temperature=1.0):
+    # Each target has its own rebased prefix. No discarded history is hidden in a cache.
+    ids=torch.cat([prefix,response[:,:-1]],1)
+    out=model(input_ids=ids,use_cache=False)
+    p=prefix.shape[1]
+    logits=out.logits[:,p-1:p+response.shape[1]-1,:visual].float()/temperature
+    return torch.log_softmax(logits,-1).gather(-1,response.unsqueeze(-1)).squeeze(-1)

code/src/futureworlds/motion_reward.py

@@ -0,0 +1,72 @@
+"""GT motion-proxy ROI and GT-referenced dynamics. No actor/search sees future GT."""
+import torch
+import torch.nn.functional as F
+from .rewards import CodecWorker
+from .batch import TensorBatch
+
+@torch.no_grad()
+def motion_regions(truth,previous):
+    """Shared across candidates: GT-derived, never shrinks after predicted disappearance."""
+    prev=torch.cat([previous[None],truth[:-1]],0)
+    motion=(truth-prev).abs().mean(1,keepdim=True)
+    thresholds=motion.flatten(1).quantile(.90,dim=1).clamp_min(.02)
+    masks=(motion>thresholds[:,None,None,None]).float()
+    masks=F.max_pool2d(masks,11,stride=1,padding=5)
+    fallback=masks.mean((1,2,3))<.01
+    masks[fallback]=1
+    crops=[];height,width=truth.shape[-2:]
+    for mask in masks:
+        y,x=torch.where(mask[0]>0)
+        y0,y1=int(y.min()),int(y.max())+1;x0,x1=int(x.min()),int(x.max())+1
+        cy=(y0+y1)//2;cx=(x0+x1)//2
+        h=max(64,y1-y0);w=max(64,x1-x0)
+        y0=max(0,min(cy-h//2,height-h));x0=max(0,min(cx-w//2,width-w))
+        crops.append((y0,min(height,y0+h),x0,min(width,x0+w)))
+    return masks,crops,fallback
+
+class MotionReward:
+    def __init__(self,codec,perceptual,scales=None,alpha=.5,beta=.25):
+        self.worker=CodecWorker(codec,perceptual);self.perceptual=perceptual
+        self.scales=scales or dict(global_scale=1.,roi_scale=1.,delta_scale=1.)
+        self.alpha=alpha;self.beta=beta
+
+    @torch.no_grad()
+    def components(self,pred,pixels):
+        truth=pixels[0,2:];previous=pixels[0,1]
+        n,h=pred.shape[:2];real=truth[None].expand(n,-1,-1,-1,-1)
+        mae=(pred-real).abs().mean((2,3,4))
+        percept=self.worker.perceptual_loss(TensorBatch(dict(real=real,pred=pred))).batch['perceptual_loss']
+        masks,crops,fallback=motion_regions(truth,previous)
+        roi_mae=((pred-real).abs()*masks).sum((2,3,4))/(3*masks.sum((1,2,3)))[None]
+        roi_percept=torch.empty((n,h),device=pred.device)
+        for t,(y0,y1,x0,x1) in enumerate(crops):
+            r=truth[t:t+1,:,y0:y1,x0:x1]
+            for start in range(0,n,2):
+                pp=pred[start:start+2,t,:,y0:y1,x0:x1]
+                roi_percept[start:start+len(pp),t]=self.perceptual(pp*2-1,r.expand(len(pp),-1,-1,-1)*2-1).flatten()
+        pred_prev=torch.cat([previous[None,None].expand(n,1,-1,-1,-1),pred[:,:-1]],1)
+        true_prev=torch.cat([previous[None],truth[:-1]],0)
+        delta=((pred-pred_prev)-(truth-true_prev)[None]).abs().mean((2,3,4))
+        return dict(mae=mae,lpips=percept,global_loss=mae+percept,
+                    roi_loss=roi_mae+roi_percept,roi_mae=roi_mae,roi_lpips=roi_percept,
+                    delta_loss=delta,roi_fallback=fallback,roi_area=masks.mean((1,2,3)))
+
+    def combine(self,c):
+        s=self.scales
+        return -(c['global_loss']+self.alpha*s['global_scale']/s['roi_scale']*c['roi_loss']+
+                 self.beta*s['global_scale']/s['delta_scale']*c['delta_loss']).mean(-1)
+
+    @torch.no_grad()
+    def decode(self,frames,context):
+        codes=torch.stack([f['response'] for f in frames],1)
+        assert codes.min()>=0 and codes.max()<4375
+        return self.worker.detokenize(TensorBatch(dict(tokens=codes,ctx_tokens=context.expand(len(codes),-1,-1))),None).batch['pixels'][:,1:]
+
+    @torch.no_grad()
+    def __call__(self,frames,context,actions,pixels):
+        c=self.components(self.decode(frames,context),pixels)
+        reward=self.combine(c)
+        assert torch.isfinite(reward).all()
+        metrics={'critic/recon_loss/mean':float(c['mae'].mean()),'critic/perceptual_loss/mean':float(c['lpips'].mean())}
+        metrics.update({k:float(v.float().mean()) for k,v in c.items()})
+        return reward,metrics

code/src/futureworlds/objective.py

@@ -0,0 +1,203 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import numpy as np
+from collections import defaultdict
+from types import SimpleNamespace
+
+def masked_mean(values, mask, axis=None):
+    """Compute mean of tensor with a masked values."""
+    return (values * mask).sum(axis=axis) / (mask.sum(axis=axis) + 1e-8)
+
+verl_F=SimpleNamespace(masked_mean=masked_mean)
+
+def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor,
+                                   response_mask: torch.Tensor,
+                                   index: np.ndarray,
+                                   epsilon: float = 1e-6,
+                                   uniform_std=False):
+    """
+    Compute advantage for GRPO, operating only on Outcome reward 
+    (with only one scalar reward for each response).
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        Returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+    """
+    scores = token_level_rewards.sum(dim=-1)
+
+    id2score = defaultdict(list)
+    id2mean = {}
+    id2std = {}
+
+    with torch.no_grad():
+        bsz = scores.shape[0]
+        for i in range(bsz):
+            id2score[index[i]].append(scores[i])
+        for idx in id2score:
+            if len(id2score[idx]) == 1:
+                id2mean[idx] = torch.tensor(0.0)
+                id2std[idx] = torch.tensor(1.0)
+            elif len(id2score[idx]) > 1:
+                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
+                id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
+            else:
+                raise ValueError(f"no score in prompt index: {idx}")
+        if uniform_std:
+            std = torch.tensor(list(id2std.values())).mean()
+            for i in range(bsz):
+                scores[i] = (scores[i] - id2mean[index[i]]) / (std + epsilon)
+        else:
+            for i in range(bsz):
+                scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
+        scores = scores.unsqueeze(-1) * response_mask
+
+    return scores, scores
+
+def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str):
+    """
+    Aggregate the loss matrix into a scalar.
+    Args:
+        loss_mat: `(torch.Tensor)`
+            shape: (bs, response_length)
+        loss_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+        loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
+            "token-mean" is the default behavior
+    Returns:
+        loss: `a scalar torch.Tensor`
+            aggregated loss
+    """
+    if loss_agg_mode == "token-mean":
+        loss = verl_F.masked_mean(loss_mat, loss_mask)
+    elif loss_agg_mode == "seq-mean-token-sum":
+        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)  # token-sum
+        loss = torch.mean(seq_losses)  # seq-mean
+    elif loss_agg_mode == "seq-mean-token-mean":
+        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1)  # token-mean
+        loss = torch.mean(seq_losses)  # seq-mean
+    else:
+        raise ValueError(f"Invalid loss_agg_mode: {loss_agg_mode}")
+
+    return loss
+
+def compute_policy_loss(old_log_prob,
+                        log_prob,
+                        advantages,
+                        response_mask,
+                        cliprange=None,
+                        cliprange_low=None,
+                        cliprange_high=None,
+                        clip_ratio_c=3.0,
+                        loss_agg_mode="token-mean"):
+    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122
+    Args:
+        old_log_prob: `(torch.Tensor)`
+            shape: (bs, response_length)
+        log_prob: `(torch.Tensor)`
+            shape: (bs, response_length)
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+        cliprange: (float)
+            The clip range used in PPO. See https://arxiv.org/abs/1707.06347
+        cliprange_low: (float)
+            The lower clip range used in PPO.
+        cliprange_high: (float)
+            The higher clip range used in PPO.
+        clip_ratio_c: (float) default: 3.0
+            The lower bound of the ratio for dual-clip PPO, See https://arxiv.org/pdf/1912.09729
+        loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
+            "token-mean" is the default behavior        
+
+    Returns:
+        pg_loss: `a scalar torch.Tensor`
+            policy gradient loss computed via PPO
+        pg_clipfrac: (float)
+            the fraction of policy gradient loss being clipped
+        ppo_kl: (float)
+            the estimated KL divergence between the latest updating policy and the old sampling policy
+        pg_clipfrac_lower: (float)
+            the fraction of policy gradient loss being clipped when the advantage is negative
+    """
+    assert clip_ratio_c > 1.0, f"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0, but get the value: {clip_ratio_c}."
+
+    negative_approx_kl = log_prob - old_log_prob
+    ratio = torch.exp(negative_approx_kl)
+    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)
+
+    pg_losses1 = -advantages * ratio
+    if cliprange_low is None:
+        cliprange_low = cliprange
+    if cliprange_high is None:
+        cliprange_high = cliprange
+    pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low,
+                                           1 + cliprange_high)  # - clip(ratio, 1-cliprange, 1+cliprange) * A
+    clip_pg_losses1 = torch.maximum(pg_losses1,
+                                    pg_losses2)  # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A)
+    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)
+
+    pg_losses3 = -advantages * clip_ratio_c
+    clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)
+    pg_clipfrac_lower = verl_F.masked_mean(
+        torch.gt(clip_pg_losses2, pg_losses3) * (advantages < 0).float(), response_mask)
+
+    pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)
+    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
+
+    return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower
+
+def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:
+    """Compute KL divergence given logprob and ref_logprob.
+    Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104
+
+    Args:
+        logprob:
+        ref_logprob:
+
+    Returns:
+
+    """
+    if kl_penalty == "kl":
+        return logprob - ref_logprob
+
+    if kl_penalty == "abs":
+        return (logprob - ref_logprob).abs()
+
+    if kl_penalty == "mse":
+        return 0.5 * (logprob - ref_logprob).square()
+
+    # J. Schulman. Approximating kl divergence, 2020.
+    # # URL http://joschu.net/blog/kl-approx.html.
+    if kl_penalty == 'low_var_kl':
+        kl = ref_logprob - logprob
+        ratio = torch.exp(kl)
+        kld = (ratio - kl - 1).contiguous()
+        return torch.clamp(kld, min=-10, max=10)
+
+    if kl_penalty == "full":
+        # so, here logprob and ref_logprob should contain the logits for every token in vocabulary
+        raise NotImplementedError
+
+    raise NotImplementedError

code/src/futureworlds/ordinary_search.py

@@ -0,0 +1,71 @@
+"""Global beam width 8; return top 4 by raw cumulative model score, no diversity penalty."""
+import torch
+from .memory import make_prefix
+
+@torch.no_grad()
+def rollout_ordinary_beam(model, context, anchor, actions, horizon, penalty,
+                       groups=4, beams=2, visual=4375, width=80):
+    assert groups == 4 and beams == 2
+    groups, beams, penalty = 1, 8, 0.0
+    returned = 4
+    model.eval()
+    device=context.device; count=groups*beams
+    history={1:anchor.expand(count,-1)}; frames=[]; traces=[]
+    scores=torch.full((count,),float('-inf'),device=device,dtype=torch.float64)
+    scores[::beams]=0
+    raw=torch.zeros(count,device=device,dtype=torch.float64)
+    for target in range(2,horizon+2):
+        original,kept=make_prefix(context.expand(count,-1),history,actions,target)
+        ancestors=torch.arange(count,device=device)
+        codes=torch.empty((count,0),device=device,dtype=torch.long)
+        logs=torch.empty((count,0),device=device)
+        pending=original; cache=None
+        for token in range(width):
+            out=model(input_ids=pending,past_key_values=cache,use_cache=True)
+            lp=out.logits[:,-1,:visual].float().log_softmax(-1)
+            parents=[]; choices=[]; next_scores=[]
+            for g in range(groups):
+                start=g*beams
+                if g:
+                    frequency=torch.bincount(torch.cat(choices),minlength=visual).double()/(g*beams)
+                else:
+                    frequency=torch.zeros(visual,device=device,dtype=torch.float64)
+                totals=scores[start:start+beams,None]+lp[start:start+beams].double()-penalty*frequency
+                score,flat=totals.flatten().topk(beams,sorted=True)
+                parents.append(start+flat//visual);choices.append(flat%visual);next_scores.append(score)
+            parent=torch.cat(parents);chosen=torch.cat(choices)
+            assert torch.equal(parent//beams,torch.arange(count,device=device)//beams)
+            raw=raw[parent]+lp[parent,chosen].double()
+            scores=torch.cat(next_scores)
+            codes=torch.cat([codes[parent],chosen[:,None]],1)
+            logs=torch.cat([logs[parent],lp[parent,chosen,None]],1)
+            ancestors=ancestors[parent]
+            cache=out.past_key_values;cache.reorder_cache(parent)
+            pending=chosen[:,None]
+        for f in frames:
+            for key in ['prefix','response','model_logp']:f[key]=f[key][ancestors]
+        history={t:x[ancestors] for t,x in history.items()}
+        history[target]=codes
+        frames.append(dict(target=target,prefix=original[ancestors],response=codes,model_logp=logs))
+        traces.append(dict(target=target,retained=kept,prefix_length=original.shape[1],
+                           frame_parent_indices=ancestors.tolist(),positions_reset=True,kv_reset=True,
+                           future_observations=False,random_sampling=False))
+    selected=torch.arange(returned,device=device)
+    frames=[{k:(v[selected] if torch.is_tensor(v) else v) for k,v in f.items()} for f in frames]
+    history={t:x[selected] for t,x in history.items()}
+    for f in frames:
+        rebuilt,_=make_prefix(context.expand(returned,-1),history,actions,f['target'])
+        assert torch.equal(rebuilt,f['prefix'])
+    rebuilt=sum(f['model_logp'].double().sum(-1) for f in frames)
+    assert torch.allclose(rebuilt,raw[selected],atol=1e-7,rtol=0)
+    codes=torch.stack([f['response'] for f in frames],1)
+    pair=[float((codes[i]!=codes[j]).float().mean()) for i in range(returned) for j in range(i+1,returned)]
+    perframe=(codes!=codes[:1]).any(0).sum(-1).tolist()
+    return frames,dict(strategy='ordinary_global_beam8_top4_full_video',groups=groups,beams_per_group=beams,
+                        active_paths=count,returned_trajectories=returned,penalty=penalty,
+                        unique_trajectories=len({tuple(c.flatten().tolist()) for c in codes}),
+                        mean_pairwise_token_difference_fraction=sum(pair)/len(pair),
+                        pairwise_token_difference_fraction=pair,different_token_positions=sum(perframe),
+                        different_positions_per_frame=perframe,model_log_scores=raw[selected].tolist(),
+                        search_scores=scores[selected].tolist(),score_is_behavior_logprob=False,
+                        on_policy_grpo=False,frames=traces)

code/src/futureworlds/reward_core.py

@@ -0,0 +1,98 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from .batch import TensorBatch as DataProto
+
+def msp_reward_fn(self, batch: DataProto, pixels, return_reward_tensor=True, save_pred=False, pixels_before_repeat=None):
+    if pixels_before_repeat is None:
+        pixels_before_repeat = pixels
+
+    batch_size = batch.batch["responses"].shape[0]
+    segment_length = self.config.data.video.segment_length
+    tokens_per_frame = self.config.processor.tokens_per_frame
+    action_dim = self.config.processor.action_dim
+
+    output_tokens = batch.batch["responses"].reshape(batch_size, segment_length - 1, tokens_per_frame + action_dim)
+    output_tokens = output_tokens[:, :, :tokens_per_frame]
+    output_tokens = output_tokens.clamp(0, self.config.processor.visual_token_num - 1).long()
+
+    ctx_tokens = batch.batch["ctx_tokens"]
+    output_tokens = output_tokens.reshape(batch_size, segment_length - 1, tokens_per_frame)
+
+    detokenize_output = self.tokenizer_wg.detokenize(
+        DataProto.from_single_dict({"tokens": output_tokens, "ctx_tokens": ctx_tokens}),
+        # DataProto.from_single_dict({"real": pixels_before_repeat[:, 2:]}, meta_info={'lpips': True}),
+        DataProto.from_single_dict({"dummy": torch.zeros((batch_size, 1))}, meta_info={'lpips': True, 'recon': self.config.trainer.reward_fn}),
+    )
+
+    if 'recon_loss' in detokenize_output.batch.keys():
+        recon_loss = detokenize_output.batch['recon_loss']
+    else:
+        pred = detokenize_output.batch['pixels'].clamp(0.0, 1.0)[:, 1:]
+        real = pixels[:, 2:]
+        if self.config.trainer.reward_fn == 'mse':
+            recon_loss = torch.mean((real - pred)**2, dim=(2, 3, 4))
+        elif self.config.trainer.reward_fn == 'mae':
+            recon_loss = torch.mean(torch.abs(real - pred), dim=(2, 3, 4))
+        else:
+            raise NotImplementedError(f"Unsupported reward function: {self.config.trainer.reward_fn}")
+
+    if 'perceptual_loss' in detokenize_output.batch.keys():
+        perceptual_loss = detokenize_output.batch['perceptual_loss']
+    else:
+        perceptual_loss = self.tokenizer_wg.perceptual_loss(DataProto.from_single_dict({"real": real, "pred": pred}))
+        perceptual_loss = perceptual_loss.batch['perceptual_loss']
+
+    if self.config.trainer.msp_reward_aggregate == "mean":
+        loss = (recon_loss + perceptual_loss).mean(-1)
+    elif self.config.trainer.msp_reward_aggregate == "last":
+        loss = (recon_loss + perceptual_loss)[:, -1]
+    elif self.config.trainer.msp_reward_aggregate == "discount":
+        discount = 0.9
+        weight = discount ** torch.arange(recon_loss.shape[1], device=recon_loss.device)
+        loss = (recon_loss + perceptual_loss) * weight.unsqueeze(0)
+        loss = loss.sum(-1) / weight.sum()
+    if not return_reward_tensor:
+        return loss
+
+    print("loss", loss.mean().item(), flush=True)
+    print("recon_loss", recon_loss.mean().item(), flush=True)
+    print("perceptual_loss", perceptual_loss.mean().item(), flush=True)
+    if save_pred:
+        pred = detokenize_output.batch['pixels'].clamp(0.0, 1.0)[:, 1:]
+        real = pixels[:, 2:]
+        for i in range(1):
+        # for i in range(5):
+            real_traj = torch.cat([real[i:i+1, j] for j in range(real.shape[1])], dim=-1)
+            pred_traj = torch.cat([pred[i:i+1, j] for j in range(pred.shape[1])], dim=-1)
+            plot_img(torch.cat([
+                real_traj.float(),
+                pred_traj.float(),
+                torch.abs(real_traj-pred_traj).float(),
+            ], dim=-2), f"{self.config.trainer.experiment_name}-{i}")
+
+    reward_tensor = torch.zeros_like(batch.batch['responses'], dtype=torch.float32)
+    for i in range(len(batch)):
+        data_item = batch[i]
+        prompt_ids = data_item.batch['prompts']
+        prompt_length = prompt_ids.shape[-1]
+        valid_response_length = data_item.batch['attention_mask'][prompt_length:].sum().long().item()
+        reward_tensor[i, valid_response_length - 1] = -loss[i].item()
+
+    return reward_tensor, {
+        "critic/recon_loss/mean": recon_loss.mean().item(),
+        "critic/perceptual_loss/mean": perceptual_loss.mean().item(),
+    }

code/src/futureworlds/rewards.py

@@ -0,0 +1,48 @@
+"""Use the upstream msp_reward_fn via a local frozen-tokenizer worker adapter."""
+from types import SimpleNamespace as NS
+import torch
+from .batch import TensorBatch
+from .reward_core import msp_reward_fn
+
+class CodecWorker:
+    def __init__(self,codec,perceptual):self.codec=codec;self.perceptual=perceptual
+    @torch.no_grad()
+    def detokenize(self,batch,dummy):
+        codes=batch.batch['tokens'];ctx=batch.batch['ctx_tokens'];movies=[]
+        for b in range(codes.shape[0]):
+            parts=[];first=None
+            for start in range(0,codes.shape[1],2):
+                decoded=self.codec.detokenize(ctx[b:b+1],codes[b:b+1,start:start+2]).float().clamp(0,1)
+                if first is None:first=decoded[:,:1]
+                parts.append(decoded[:,1:])
+            movies.append(torch.cat([first]+parts,1))
+        return TensorBatch({'pixels':torch.cat(movies,0)})
+    @torch.no_grad()
+    def perceptual_loss(self,batch):
+        real,pred=batch.batch['real'],batch.batch['pred'];values=[]
+        for r,p in zip(real,pred):
+            pieces=[]
+            for i in range(0,len(p),2):pieces.append(self.perceptual(p[i:i+2]*2-1,r[i:i+2]*2-1).flatten())
+            values.append(torch.cat(pieces))
+        return TensorBatch({'perceptual_loss':torch.stack(values)})
+
+class VideoReward:
+    msp_reward_fn=msp_reward_fn
+    def __init__(self,codec,perceptual,horizon):
+        self.config=NS(data=NS(video=NS(segment_length=horizon+1)),processor=NS(tokens_per_frame=80,action_dim=13,visual_token_num=4375),trainer=NS(reward_fn='mae',msp_reward_aggregate='mean'))
+        self.tokenizer_wg=CodecWorker(codec,perceptual)
+    @torch.no_grad()
+    def __call__(self,frames,context,actions,pixels):
+        # Only reward bookkeeping uses concatenated responses. Policy scoring uses actual per-frame prefixes.
+        codes=torch.stack([f['response'] for f in frames],1)
+        assert codes.min()>=0 and codes.max()<4375,'Illegal generated token; never hide it with upstream clamp'
+        group,horizon,_=codes.shape
+        act=actions[2:2+horizon].unsqueeze(0).expand(group,-1,-1)
+        responses=torch.cat([codes,act],-1).reshape(group,-1)
+        prompt=frames[0]['prefix']
+        batch=TensorBatch({'responses':responses,'prompts':prompt,'attention_mask':torch.ones(group,prompt.shape[1]+responses.shape[1],device=codes.device,dtype=torch.long),'ctx_tokens':context.expand(group,-1,-1)})
+        truth=pixels.expand(group,-1,-1,-1,-1)
+        tensor,metrics=self.msp_reward_fn(batch,truth)
+        reward=tensor.sum(-1)
+        assert torch.isfinite(reward).all()
+        return reward,metrics

code/src/futureworlds/search.py

@@ -0,0 +1,72 @@
+"""Four disjoint beam groups, with unmodified policy scores for the RL surrogate.
+
+Search penalties only affect proposal selection. Deterministic proposals do not
+have on-policy sampling probabilities; the beam RL objective remains biased.
+"""
+import torch
+from .memory import make_prefix
+
+@torch.no_grad()
+def rollout_group_beam(model, context, anchor, actions, horizon, penalty,
+                       groups=4, beams=2, visual=4375, width=80):
+    model.eval()
+    device=context.device; count=groups*beams
+    history={1:anchor.expand(count,-1)}; frames=[]; traces=[]
+    scores=torch.full((count,),float('-inf'),device=device,dtype=torch.float64)
+    scores[::beams]=0
+    raw=torch.zeros(count,device=device,dtype=torch.float64)
+    for target in range(2,horizon+2):
+        original,kept=make_prefix(context.expand(count,-1),history,actions,target)
+        ancestors=torch.arange(count,device=device)
+        codes=torch.empty((count,0),device=device,dtype=torch.long)
+        logs=torch.empty((count,0),device=device)
+        pending=original; cache=None
+        for token in range(width):
+            out=model(input_ids=pending,past_key_values=cache,use_cache=True)
+            lp=out.logits[:,-1,:visual].float().log_softmax(-1)
+            parents=[]; choices=[]; next_scores=[]
+            for g in range(groups):
+                start=g*beams
+                if g:
+                    frequency=torch.bincount(torch.cat(choices),minlength=visual).double()/(g*beams)
+                else:
+                    frequency=torch.zeros(visual,device=device,dtype=torch.float64)
+                totals=scores[start:start+beams,None]+lp[start:start+beams].double()-penalty*frequency
+                score,flat=totals.flatten().topk(beams,sorted=True)
+                parents.append(start+flat//visual);choices.append(flat%visual);next_scores.append(score)
+            parent=torch.cat(parents);chosen=torch.cat(choices)
+            assert torch.equal(parent//beams,torch.arange(count,device=device)//beams)
+            raw=raw[parent]+lp[parent,chosen].double()
+            scores=torch.cat(next_scores)
+            codes=torch.cat([codes[parent],chosen[:,None]],1)
+            logs=torch.cat([logs[parent],lp[parent,chosen,None]],1)
+            ancestors=ancestors[parent]
+            cache=out.past_key_values;cache.reorder_cache(parent)
+            pending=chosen[:,None]
+        for f in frames:
+            for key in ['prefix','response','model_logp']:f[key]=f[key][ancestors]
+        history={t:x[ancestors] for t,x in history.items()}
+        history[target]=codes
+        frames.append(dict(target=target,prefix=original[ancestors],response=codes,model_logp=logs))
+        traces.append(dict(target=target,retained=kept,prefix_length=original.shape[1],
+                           frame_parent_indices=ancestors.tolist(),positions_reset=True,kv_reset=True,
+                           future_observations=False,random_sampling=False))
+    selected=torch.arange(0,count,beams,device=device)
+    frames=[{k:(v[selected] if torch.is_tensor(v) else v) for k,v in f.items()} for f in frames]
+    history={t:x[selected] for t,x in history.items()}
+    for f in frames:
+        rebuilt,_=make_prefix(context.expand(groups,-1),history,actions,f['target'])
+        assert torch.equal(rebuilt,f['prefix'])
+    rebuilt=sum(f['model_logp'].double().sum(-1) for f in frames)
+    assert torch.allclose(rebuilt,raw[selected],atol=1e-7,rtol=0)
+    codes=torch.stack([f['response'] for f in frames],1)
+    pair=[float((codes[i]!=codes[j]).float().mean()) for i in range(groups) for j in range(i+1,groups)]
+    perframe=(codes!=codes[:1]).any(0).sum(-1).tolist()
+    return frames,dict(strategy='grouped_diverse_full_video_beam',groups=groups,beams_per_group=beams,
+                        active_paths=count,returned_trajectories=groups,penalty=penalty,
+                        unique_trajectories=len({tuple(c.flatten().tolist()) for c in codes}),
+                        mean_pairwise_token_difference_fraction=sum(pair)/len(pair),
+                        pairwise_token_difference_fraction=pair,different_token_positions=sum(perframe),
+                        different_positions_per_frame=perframe,model_log_scores=raw[selected].tolist(),
+                        search_scores=scores[selected].tolist(),score_is_behavior_logprob=False,
+                        on_policy_grpo=False,frames=traces)

code/src/futureworlds/text_conditioning.py

@@ -0,0 +1,73 @@
+"""Pretrained-text conditioning with zero-initialized residual cross-attention.
+
+Frozen T5 features are inputs, not vocabulary IDs. Visual/action positions and
+vocabulary stay unchanged. No-text requests bypass the added residual exactly.
+"""
+import torch
+_UNSET = object()
+from torch import nn
+import torch.nn.functional as F
+
+
+class TextResidual(nn.Module):
+    def __init__(self, hidden=768, text_width=768, inner=256, heads=4):
+        super().__init__();self.heads=heads;self.inner=inner
+        self.query_norm=nn.LayerNorm(hidden)
+        self.text_norm=nn.LayerNorm(text_width)
+        self.q=nn.Linear(hidden,inner,bias=False)
+        self.k=nn.Linear(text_width,inner,bias=False)
+        self.v=nn.Linear(text_width,inner,bias=False)
+        self.out=nn.Linear(inner,hidden,bias=False)
+        nn.init.zeros_(self.out.weight)
+    def forward(self, hidden, text, mask):
+        b,t,_=hidden.shape
+        if text.shape[0]!=b:
+            assert text.shape[0]==1, 'Generation expansion is valid only within a single episode'
+            text=text.expand(b,-1,-1);mask=mask.expand(b,-1)
+        dim=self.inner//self.heads
+        q=self.q(self.query_norm(hidden)).view(b,t,self.heads,dim).transpose(1,2)
+        memory=self.text_norm(text)
+        k=self.k(memory).view(b,-1,self.heads,dim).transpose(1,2)
+        v=self.v(memory).view(b,-1,self.heads,dim).transpose(1,2)
+        value=F.scaled_dot_product_attention(q,k,v,attn_mask=mask[:,None,None,:].bool(),dropout_p=0.)
+        value=value.transpose(1,2).reshape(b,t,self.inner)
+        # Explicitly suppress all-padding / deliberately dropped instruction.
+        value=self.out(value)*mask.any(-1)[:,None,None].to(value.dtype)
+        return hidden+value.to(hidden.dtype)
+
+
+class TextWorldModel(nn.Module):
+    def __init__(self, core, text_width=768, layers=(3,7,11), inner=256):
+        super().__init__();self.core=core;self.layers=list(layers)
+        self.adapters=nn.ModuleDict({str(i):TextResidual(core.config.hidden_size,text_width,inner) for i in layers})
+        self._text=None;self._mask=None;self._handles=[]
+        for i in layers:
+            self._handles.append(core.model.layers[i].register_forward_pre_hook(self._hook(i),with_kwargs=True))
+    def _hook(self,index):
+        def hook(module,args,kwargs):
+            if self._text is None:return args,kwargs
+            adapter=self.adapters[str(index)]
+            if args:
+                return (adapter(args[0],self._text,self._mask),)+args[1:],kwargs
+            kw=dict(kwargs);kw['hidden_states']=adapter(kw['hidden_states'],self._text,self._mask)
+            return args,kw
+        return hook
+    def set_text(self,text,mask):
+        assert (text is None)==(mask is None)
+        if text is not None:
+            assert text.ndim==3 and mask.shape==text.shape[:2]
+            assert torch.isfinite(text).all()
+        self._text=text;self._mask=mask
+    def forward(self,text_features=_UNSET,text_mask=_UNSET,**kwargs):
+        # Generation and selected-token scoring use the explicitly bound episode.
+        # No activation checkpointing: mutable conditioning cannot change during replay.
+        if text_features is not _UNSET:
+            assert text_mask is not _UNSET
+            self.set_text(text_features,text_mask)
+        else:
+            assert text_mask is _UNSET
+        return self.core(**kwargs)
+    def generate(self,*args,**kwargs):
+        # Caller sets the same episode instruction before every generated frame.
+        return self.core.generate(*args,**kwargs)
+    def compute_transition_scores(self,*args,**kwargs):return self.core.compute_transition_scores(*args,**kwargs)

robocasa/README.md

@@ -0,0 +1,50 @@
+---
+language:
+- en
+tags:
+- futureworlds
+- world-model
+- robotics
+- video-prediction
+- reinforcement-learning
+---
+# FutureWorlds — RoboCasa
+
+Local release candidate, not yet published. Model weights are verified by CPU SHA-256; end-to-end inference with the portable RGB pipeline is not yet validated.
+
+## Checkpoint variants
+
+| Folder | Method | Updates | Use |
+|---|---|---:|---|
+| `models/main` | MemSPO (R0) | 400 | Paper main table |
+| `models/sft` | SFT | 50,000 | Initialization/control |
+| `models/memspo200` | MemSPO (R0) | 200 | Matched ablation |
+| `models/grpo200` | GRPO (R0) | 200 | Matched ablation |
+| `models/ordinary200` | Ordinary global beam8, top4 (R0) | 200 | Proposal-strategy ablation |
+
+All post-training checkpoints start from their paired SFT checkpoint. `main` and `memspo200` must not be interchanged.
+
+## Main-table evaluation
+
+128 trajectories, 32 predicted frames, 2 initial observations, per-frame Beam4, anchor1 + recent6, no true-future refresh. Recorded split: held-out test; runtime: GPU FP32.
+
+| PSNR ↑ | SSIM ↑ | LPIPS ↓ |
+|---:|---:|---:|
+| 18.4719 | 0.765948 | 0.183338 |
+
+These are recorded paper results; staging weights does not constitute a fresh reproduction.
+
+## Required components
+
+- Dataset-specific `codec/` and `preprocessing/` files, paired by the manifest.
+- Trained text adapters in each checkpoint and frozen T5-base at revision `a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1`; T5 is referenced, not duplicated.
+- 52 action tokens per predicted frame. RoboCasa uses four ordered 13-dimensional control blocks; do not flatten time in a different order.
+- Original FP32 safetensors; no post-hoc quantization or checkpoint merging.
+
+## Loading and scope
+
+Use `futureworlds.checkpoints.load_world_model(bundle_root, "robocasa", variant="main")` for the predictor only. This does not decode RGB, prepare actions, or download T5 automatically. A complete RGB demo and clean-environment GPU equivalence run remain pending. These models are custom action-conditioned predictors, not generic `pipeline("text-to-video")` models.
+
+## Provenance and limitations
+
+Core architecture and tokenizer derive from RLVR-World/iVideoGPT. Upstream license texts are included under `licenses/`; the license for original additions is not yet selected. The raw datasets and third-party baseline weights are not included. Learned-world feedback examples are qualitative generated trajectories, not independently verified physical task-success results.

robocasa/codec/config.json

@@ -0,0 +1,40 @@
+{
+  "_class_name": "CompressiveVQModelFSQ",
+  "_diffusers_version": "0.35.2",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    256,
+    512
+  ],
+  "context_length": 1,
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "dyn_fsq_levels": 12,
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 64,
+  "layers_per_block": 2,
+  "lookup_from_codebook": true,
+  "max_att_resolution": 32,
+  "mid_block_add_attention": false,
+  "norm_num_groups": 32,
+  "norm_type": "group",
+  "out_channels": 3,
+  "patch_size": 4,
+  "resolution": 256,
+  "sample_size": 32,
+  "scaling_factor": 0.18215,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "vq_fsq_levels": 12
+}

robocasa/codec/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ea6133eae9d3deba4fc546ac7d0cb318363d0d79372b29a8e7d1a753023b77d
+size 517850360

robocasa/futureworlds_config.json

@@ -0,0 +1,77 @@
+{
+  "schema_version": 1,
+  "dataset": "robocasa",
+  "status": "staged_research_preview",
+  "codec": "robocasa/codec",
+  "action_ranges": "robocasa/preprocessing/action_ranges.pt",
+  "action_tokens_per_step": 52,
+  "text_conditioning": {
+    "enabled": true,
+    "repo_id": "google-t5/t5-base",
+    "revision": "a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1",
+    "layers": [
+      3,
+      7,
+      11
+    ],
+    "inner": 256,
+    "width": 768,
+    "max_tokens": 128,
+    "bundled": false
+  },
+  "initial_observations": 2,
+  "memory": {
+    "anchor_states": 1,
+    "recent_states": 6,
+    "ground_truth_refresh": false
+  },
+  "evaluation": {
+    "strategy": "per_frame_beam",
+    "beam_width": 4,
+    "horizon": 32,
+    "visual_tokens_per_frame": 80,
+    "valid_visual_vocab": 4375
+  },
+  "variants": {
+    "main": {
+      "path": "robocasa/models/main",
+      "method": "MemSPO",
+      "updates": 400,
+      "reward": "R0",
+      "purpose": "paper_main_table",
+      "text": true
+    },
+    "sft": {
+      "path": "robocasa/models/sft",
+      "method": "SFT",
+      "updates": 50000,
+      "reward": null,
+      "purpose": "matched_ablation",
+      "text": true
+    },
+    "memspo200": {
+      "path": "robocasa/models/memspo200",
+      "method": "MemSPO",
+      "updates": 200,
+      "reward": "R0",
+      "purpose": "matched_ablation",
+      "text": true
+    },
+    "grpo200": {
+      "path": "robocasa/models/grpo200",
+      "method": "GRPO",
+      "updates": 200,
+      "reward": "R0",
+      "purpose": "matched_ablation",
+      "text": true
+    },
+    "ordinary200": {
+      "path": "robocasa/models/ordinary200",
+      "method": "OrdinaryBeam",
+      "updates": 200,
+      "reward": "R0",
+      "purpose": "matched_ablation",
+      "text": true
+    }
+  }
+}

robocasa/models/grpo200/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

robocasa/models/grpo200/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

robocasa/models/grpo200/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af36ba44677963a0de2753dcf933c46f616505792fe4d6668138adc2d8e59720
+size 508419080

robocasa/models/grpo200/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eba7ba8715d82c79f8909a8d65cf3ede5f1874b0963efcd3a351f2f8ca13633c
+size 9476032

robocasa/models/main/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

robocasa/models/main/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

robocasa/models/main/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0d4420b480f44f9f1f970edbfc96a95adb5e428dc6b49ca3d1b2c61fddb5807
+size 508419080

robocasa/models/main/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92ad8666771696ee22f6c7aaee12d9bd8bd6195569298defb1b89b3bcba11d1e
+size 9476032

robocasa/models/memspo200/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

robocasa/models/memspo200/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

robocasa/models/memspo200/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6193281a747b6155cb32b7e4e3844f1ff05ac0e07d9df4cd3f11bba4c8711d92
+size 508419080

robocasa/models/memspo200/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c9f8ac2cc6c128fc29305546a29effbd99250f47905e7711a8330a397bfd7c0
+size 9476032

robocasa/models/ordinary200/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

robocasa/models/ordinary200/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

robocasa/models/ordinary200/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:486a6c324466839121a4fb392450b1093e98d7482dbb9b6c794baf4ebb7b1519
+size 508419080

robocasa/models/ordinary200/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4fc6d76fe21e5687fa7dd4cb2c82da0b3c52554490be6211e22aa4fabc85900
+size 9476032

robocasa/models/sft/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 9006,
+  "dtype": "float32",
+  "eos_token_id": 9007,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_length": null,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 9008
+}

robocasa/models/sft/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 9006,
+  "eos_token_id": 9007,
+  "transformers_version": "4.57.1"
+}

robocasa/models/sft/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a86d465e8b5c1cc824238b010aed477ea61d6d8307102135b886b32ca065463c
+size 508419080

robocasa/models/sft/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef8a06e8ba6451ad1873d9476400c04ae581ce8e654ca77958f7cc56d9f6b380
+size 9476032

robocasa/preprocessing/action_ranges.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55531c0e5296d01b8eb78aa941b422bde30e80e458a24d362f874671f2dbfaf4
+size 2003

robocasa/training_recipe.json

@@ -0,0 +1,23 @@
+{
+  "max_text_tokens": 128,
+  "text_layers": [
+    3,
+    7,
+    11
+  ],
+  "text_inner": 256,
+  "seed": 20260918,
+  "sft_steps": 50000,
+  "sft_global_batch": 64,
+  "core_lr": 1e-05,
+  "adapter_lr": 0.0001,
+  "warmup_steps": 500,
+  "rl_steps": 500,
+  "rl_lr": 1e-06,
+  "kl_coef": 0.001,
+  "clip": 0.2,
+  "horizon": 10,
+  "groups": 4,
+  "beams_per_group": 2,
+  "diversity_penalty": 0.1
+}

rt1/models/main/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c426c478c690679230b613401b59f8d17bc7863b33450ba12964d42722b489cd
+size 9476032

rt1/models/sft/text_adapters.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:006511fd5bbfdeb2f6975695415243b5cec4912517c7f187e65864a48ecfcf07
+size 9476032