Upload all Matrix-Game evaluation results

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/case0_full_original.json

@@ -0,0 +1,81 @@
+{
+  "case": "case0_full_original",
+  "raw_rgb_psnr": [
+    25.77519416809082,
+    4.8833327293396,
+    4.636667251586914,
+    5.688905715942383,
+    7.4164934158325195,
+    7.8217339515686035,
+    8.03274154663086,
+    8.007763862609863,
+    8.02673625946045,
+    8.09578800201416,
+    8.218385696411133,
+    8.326114654541016,
+    7.63754940032959,
+    7.489852428436279,
+    7.3153157234191895,
+    7.299378395080566,
+    7.418322563171387,
+    7.375185012817383,
+    7.256275177001953,
+    7.181506156921387,
+    7.253871917724609,
+    7.2649149894714355,
+    7.2623491287231445,
+    7.290315628051758
+  ],
+  "raw_rgb_mse": [
+    0.002645334927365184,
+    0.32483792304992676,
+    0.3438216745853424,
+    0.2698419392108917,
+    0.18128031492233276,
+    0.16513024270534515,
+    0.1572989672422409,
+    0.1582062542438507,
+    0.15751661360263824,
+    0.1550319343805313,
+    0.15071670711040497,
+    0.14702408015727997,
+    0.17228403687477112,
+    0.178243950009346,
+    0.18555320799350739,
+    0.18623535335063934,
+    0.18120399117469788,
+    0.18301282823085785,
+    0.1880929172039032,
+    0.19135920703411102,
+    0.1881970316171646,
+    0.18771910667419434,
+    0.18783006072044373,
+    0.18662439286708832
+  ],
+  "mean_rgb": [
+    0.6862332224845886,
+    0.1697441190481186,
+    0.156618133187294,
+    0.2281637042760849,
+    0.31492793560028076,
+    0.33408209681510925,
+    0.34403756260871887,
+    0.3432576060295105,
+    0.3432428538799286,
+    0.34543678164482117,
+    0.35083115100860596,
+    0.357102632522583,
+    0.3239423930644989,
+    0.3186531066894531,
+    0.3103964030742645,
+    0.3092810809612274,
+    0.3162357807159424,
+    0.31385642290115356,
+    0.30743733048439026,
+    0.3035033941268921,
+    0.3069319725036621,
+    0.3079564869403839,
+    0.30708178877830505,
+    0.30761298537254333
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/case1_full_original.json

@@ -0,0 +1,81 @@
+{
+  "case": "case1_full_original",
+  "raw_rgb_psnr": [
+    20.75506591796875,
+    5.314362525939941,
+    4.540973663330078,
+    5.3272600173950195,
+    5.330643653869629,
+    5.8266520500183105,
+    5.940614223480225,
+    6.0602803230285645,
+    6.2254180908203125,
+    6.322708606719971,
+    6.403367042541504,
+    6.356753826141357,
+    6.32526159286499,
+    6.268328666687012,
+    6.2384490966796875,
+    6.203701496124268,
+    6.242549896240234,
+    6.298856735229492,
+    6.294210433959961,
+    6.316972732543945,
+    6.402440547943115,
+    6.451478481292725,
+    6.471812725067139,
+    6.447417736053467
+  ],
+  "raw_rgb_mse": [
+    0.00840414222329855,
+    0.2941465377807617,
+    0.3514816462993622,
+    0.29327428340911865,
+    0.29304590821266174,
+    0.26141756772994995,
+    0.25464698672294617,
+    0.2477262318134308,
+    0.23848342895507812,
+    0.2332003265619278,
+    0.22890923917293549,
+    0.23137935996055603,
+    0.23306328058242798,
+    0.23613868653774261,
+    0.23776891827583313,
+    0.23967893421649933,
+    0.23754450678825378,
+    0.23448461294174194,
+    0.23473559319972992,
+    0.23350851237773895,
+    0.22895805537700653,
+    0.22638735175132751,
+    0.22532984614372253,
+    0.22659912705421448
+  ],
+  "mean_rgb": [
+    0.6869287490844727,
+    0.22597603499889374,
+    0.16900615394115448,
+    0.21825699508190155,
+    0.21941502392292023,
+    0.248577281832695,
+    0.2554280757904053,
+    0.26183179020881653,
+    0.27177396416664124,
+    0.27722224593162537,
+    0.281747430562973,
+    0.27876004576683044,
+    0.27639126777648926,
+    0.2728942036628723,
+    0.2711937129497528,
+    0.2694670855998993,
+    0.2724168300628662,
+    0.27583178877830505,
+    0.2756144106388092,
+    0.27678048610687256,
+    0.2816024720668793,
+    0.284187912940979,
+    0.2853650152683258,
+    0.283825546503067
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/complete.json

@@ -0,0 +1 @@
+{"mode": "full_original", "cases": 2, "frames": 24}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/launcher.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-diagnostic
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=01:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(wan_oracle full_original full_weighted i2v_original i2v_weighted)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/diagnose_matrix_game2.py" "$OUTPUT/probe.py"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/diagnose_matrix_game2.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/manifest.json

@@ -0,0 +1,59 @@
+{
+  "diagnostic_only": true,
+  "mode": "full_original",
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "weight_source": "released raw distilled",
+  "optimizer": null,
+  "lr": null,
+  "lora": null,
+  "activation_checkpointing": false,
+  "batch_size": 1,
+  "trainable_parameters": 0,
+  "denoising_steps": 3,
+  "shift": 5,
+  "local_attention_latents": 6,
+  "selector": "native sliding cache",
+  "reference": "raw RGB resized352x640, not benchmark Oasis references",
+  "targets": [
+    700,
+    724
+  ],
+  "full_history": [
+    100,
+    700
+  ],
+  "full_prefix_latents": 153,
+  "oracle": "9copies+RGB100:724 encodedonce;159latents; discard153latent prefix",
+  "i2v": {
+    "anchor": 699,
+    "decoded_slice": [
+      1,
+      25
+    ],
+    "controls": [
+      700,
+      724
+    ],
+    "right_noops": 9
+  },
+  "normalization_change": "weighted modes apply scales and accumulate RMS in fp32",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "parameter_counts": {
+    "generator_total": 1619242304,
+    "generator_trainable": 0,
+    "vae_total": 126892531,
+    "clip_total": 1193014785
+  },
+  "action_norm_instances": 60
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85090"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_original_job85088_task1/probe.py

@@ -0,0 +1,113 @@
+"""Frozen Matrix diagnostics; raw-RGB PSNR is separate from benchmark scoring."""
+import argparse
+import json
+import sys
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(ROOT / 'shared'), str(ROOT / 'Matrix-Game/Matrix-Game-2')]
+from protocol import load_minecraft_case, select_cases
+from worldmem_adapter import MatrixGameAdapter, map_actions
+from demo_utils.constant import ZERO_VAE_CACHE
+
+MODES = ['wan_oracle', 'full_original', 'full_weighted', 'i2v_original', 'i2v_weighted']
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('--mode', choices=MODES, required=True)
+parser.add_argument('--data-root', type=Path, required=True)
+parser.add_argument('--output', type=Path, required=True)
+parser.add_argument('--dry-run', action='store_true')
+
+def weighted_norm(self, x):
+    y = x.float()
+    return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight
+
+def save_result(output, name, prediction, reference):
+    prediction = prediction.float().cpu()
+    if prediction.shape != reference.shape or not torch.isfinite(prediction).all():
+        raise ValueError(f'Invalid Matrix prediction: {name}, shape={tuple(prediction.shape)}')
+    prediction = prediction.clamp(0, 1)
+    mse = (prediction - reference).square().mean((1, 2, 3))
+    report = {'case': name, 'raw_rgb_psnr': (-10 * mse.clamp_min(1e-12).log10()).tolist(),
+              'raw_rgb_mse': mse.tolist(), 'mean_rgb': prediction.mean((1, 2, 3)).tolist()}
+    (output / f'{name}.json').write_text(json.dumps(report, indent=2))
+    grid = torch.cat([torch.cat([v[t] for t in [0, 1, 2, 5, 11, 23]], dim=2)
+                      for v in (reference, prediction)], dim=1)
+    Image.fromarray(grid.mul(255).byte().permute(1, 2, 0).numpy()).save(output / f'{name}.jpg')
+    print(json.dumps({'case': name, 'mean_psnr': sum(report['raw_rgb_psnr']) / 24,
+                      'first_psnr': report['raw_rgb_psnr'][:3]}), flush=True)
+
+@torch.inference_mode()
+def main():
+    args = parser.parse_args()
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    manifest = {'diagnostic_only': True, 'mode': args.mode, 'checkpoint': str(checkpoint),
+                'weight_source': 'released raw distilled', 'optimizer': None, 'lr': None, 'lora': None,
+                'activation_checkpointing': False, 'batch_size': 1, 'trainable_parameters': 0,
+                'denoising_steps': 3, 'shift': 5, 'local_attention_latents': 6,
+                'selector': 'native sliding cache', 'reference': 'raw RGB resized352x640, not benchmark Oasis references',
+                'targets': [700, 724], 'full_history': [100, 700], 'full_prefix_latents': 153,
+                'oracle': '9copies+RGB100:724 encodedonce;159latents; discard153latent prefix',
+                'i2v': {'anchor': 699, 'decoded_slice': [1, 25], 'controls': [700, 724], 'right_noops': 9},
+                'normalization_change': 'weighted modes apply scales and accumulate RMS in fp32',
+                'cases': [{'index': r['global_test_index'], 'path': r['relative_path'],
+                           'seed': 42000000 + int(r['global_test_index'])} for r in rows]}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    from wan.modules.action_module import WanRMSNorm
+    adapter = MatrixGameAdapter(checkpoint)
+    manifest['parameter_counts'] = adapter.parameter_counts
+    manifest['action_norm_instances'] = sum(isinstance(m, WanRMSNorm) for m in adapter.pipeline.modules())
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    if args.mode.endswith('weighted'):
+        WanRMSNorm.forward = weighted_norm
+    for row in rows:
+        row.update(target_end='724', target_frames='24')
+        case = load_minecraft_case(row, args.data_root)
+        reference = F.interpolate(case['rgb'][600:624], (352, 640), mode='bilinear', align_corners=False, antialias=True)
+        seed = 42000000 + int(row['global_test_index'])
+        if args.mode == 'wan_oracle':
+            rgb = F.interpolate(case['rgb'].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            rgb = torch.cat([rgb[:1].expand(9, -1, -1, -1), rgb])
+            latent = adapter.vae.encode(rgb.transpose(0, 1)[None].mul(2).sub(1).bfloat16(), device=adapter.device, tiled=False)
+            cache, chunks = [None] * len(ZERO_VAE_CACHE), []
+            for start in range(0, latent.shape[2], 3):
+                decoded, cache = adapter.pipeline.vae_decoder(latent[:, :, start:start + 3].transpose(1, 2).half(), *cache)
+                if start >= 153:
+                    chunks.append(decoded.cpu())
+            prediction = torch.cat(chunks, dim=1)[0].float().add(1).div(2)
+            del latent, rgb, cache, decoded, chunks
+        elif args.mode.startswith('full'):
+            prediction, _ = adapter.generate(case, target_frames=24, seed=seed)
+        else:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            anchor = F.interpolate(case['rgb'][599:600].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            anchor = anchor.transpose(0, 1)[None].mul(2).sub(1).bfloat16()
+            image = anchor.new_zeros((1, 3, 33, 352, 640))
+            image[:, :, :1] = anchor
+            image_latents = adapter.vae.encode(image, device=adapter.device, tiled=False)
+            mask = torch.zeros_like(image_latents[:, :4])
+            mask[:, :, :1] = 1
+            actions = torch.cat([case['actions'][600:624], torch.zeros(9, 25)]).cuda().bfloat16()
+            keyboard, mouse = map_actions(actions, 1.0)
+            condition = {'cond_concat': torch.cat([mask, image_latents], dim=1),
+                         'visual_context': adapter.vae.clip.encode_video(anchor).bfloat16(),
+                         'keyboard_cond': keyboard[None], 'mouse_cond': mouse[None]}
+            noise = torch.randn((1, 16, 9, 44, 80), device='cuda', dtype=torch.bfloat16)
+            videos = adapter.pipeline.inference(noise, condition, initial_latent=None, mode='universal')
+            prediction = torch.cat(videos, dim=1)[0, 1:25].float().add(1).div(2)
+            del image, anchor, image_latents, condition, videos, noise
+        save_result(args.output, f"case{row['global_test_index']}_{args.mode}", prediction, reference)
+        del prediction, reference, case
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'cases': 2, 'frames': 24}))
+    print('DIAGNOSTICS_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/case0_full_weighted.json

@@ -0,0 +1,81 @@
+{
+  "case": "case0_full_weighted",
+  "raw_rgb_psnr": [
+    25.945024490356445,
+    4.916481018066406,
+    4.617745876312256,
+    5.818724632263184,
+    7.579941272735596,
+    7.90878963470459,
+    8.117096900939941,
+    8.039791107177734,
+    8.028121948242188,
+    8.168718338012695,
+    8.279888153076172,
+    8.406937599182129,
+    7.733737945556641,
+    7.428074359893799,
+    7.25970983505249,
+    7.244600296020508,
+    7.345211505889893,
+    7.292293548583984,
+    7.148654937744141,
+    7.099469184875488,
+    7.157524108886719,
+    7.194796085357666,
+    7.167351245880127,
+    7.1892571449279785
+  ],
+  "raw_rgb_mse": [
+    0.0025438847951591015,
+    0.32236796617507935,
+    0.34532293677330017,
+    0.26189520955085754,
+    0.17458458244800568,
+    0.1618531197309494,
+    0.1542731374502182,
+    0.15704382956027985,
+    0.1574663519859314,
+    0.15245023369789124,
+    0.14859738945960999,
+    0.1443132609128952,
+    0.16851019859313965,
+    0.1807975471019745,
+    0.187944233417511,
+    0.18859927356243134,
+    0.18428027629852295,
+    0.1865394413471222,
+    0.19281220436096191,
+    0.1950082778930664,
+    0.19241884350776672,
+    0.19077451527118683,
+    0.19198393821716309,
+    0.1910180002450943
+  ],
+  "mean_rgb": [
+    0.6870425343513489,
+    0.1719239205121994,
+    0.1551111340522766,
+    0.23643054068088531,
+    0.32379114627838135,
+    0.3390358090400696,
+    0.3493805527687073,
+    0.3448195159435272,
+    0.3432111144065857,
+    0.349094033241272,
+    0.3539511561393738,
+    0.360940158367157,
+    0.32845252752304077,
+    0.3160686194896698,
+    0.3070259988307953,
+    0.30618816614151,
+    0.3123764097690582,
+    0.3096041679382324,
+    0.30161795020103455,
+    0.29898521304130554,
+    0.30154550075531006,
+    0.3034396171569824,
+    0.3025643229484558,
+    0.302489310503006
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/case1_full_weighted.json

@@ -0,0 +1,81 @@
+{
+  "case": "case1_full_weighted",
+  "raw_rgb_psnr": [
+    20.358261108398438,
+    5.021709442138672,
+    4.7442827224731445,
+    5.203425407409668,
+    5.046121597290039,
+    5.245299339294434,
+    5.384633541107178,
+    5.533613681793213,
+    5.719480514526367,
+    5.826648712158203,
+    5.863295078277588,
+    5.830295085906982,
+    5.900106906890869,
+    5.944562911987305,
+    6.028284072875977,
+    5.974111080169678,
+    6.014407157897949,
+    6.091710090637207,
+    6.0647077560424805,
+    6.101002216339111,
+    6.22338342666626,
+    6.275542259216309,
+    6.29188346862793,
+    6.269043922424316
+  ],
+  "raw_rgb_mse": [
+    0.009208180010318756,
+    0.3146509528160095,
+    0.33540669083595276,
+    0.3017570674419403,
+    0.3128872513771057,
+    0.298861563205719,
+    0.28942540287971497,
+    0.27966535091400146,
+    0.2679488956928253,
+    0.26141780614852905,
+    0.2592211663722992,
+    0.26119837164878845,
+    0.25703325867652893,
+    0.2544155716896057,
+    0.24955807626247406,
+    0.25269049406051636,
+    0.2503567337989807,
+    0.24593989551067352,
+    0.2474737912416458,
+    0.24541424214839935,
+    0.2385951578617096,
+    0.23574677109718323,
+    0.23486140370368958,
+    0.23609979450702667
+  ],
+  "mean_rgb": [
+    0.6835651397705078,
+    0.20514607429504395,
+    0.18633019924163818,
+    0.21573315560817719,
+    0.20562198758125305,
+    0.21696747839450836,
+    0.2257290482521057,
+    0.23421454429626465,
+    0.24657680094242096,
+    0.2528727352619171,
+    0.25491902232170105,
+    0.25266510248184204,
+    0.25558704137802124,
+    0.25767818093299866,
+    0.26240184903144836,
+    0.25888851284980774,
+    0.2621479630470276,
+    0.26704472303390503,
+    0.2655029892921448,
+    0.267596572637558,
+    0.2748774290084839,
+    0.27775564789772034,
+    0.2788831293582916,
+    0.2775684595108032
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/complete.json

@@ -0,0 +1 @@
+{"mode": "full_weighted", "cases": 2, "frames": 24}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/launcher.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-diagnostic
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=01:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(wan_oracle full_original full_weighted i2v_original i2v_weighted)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/diagnose_matrix_game2.py" "$OUTPUT/probe.py"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/diagnose_matrix_game2.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/manifest.json

@@ -0,0 +1,59 @@
+{
+  "diagnostic_only": true,
+  "mode": "full_weighted",
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "weight_source": "released raw distilled",
+  "optimizer": null,
+  "lr": null,
+  "lora": null,
+  "activation_checkpointing": false,
+  "batch_size": 1,
+  "trainable_parameters": 0,
+  "denoising_steps": 3,
+  "shift": 5,
+  "local_attention_latents": 6,
+  "selector": "native sliding cache",
+  "reference": "raw RGB resized352x640, not benchmark Oasis references",
+  "targets": [
+    700,
+    724
+  ],
+  "full_history": [
+    100,
+    700
+  ],
+  "full_prefix_latents": 153,
+  "oracle": "9copies+RGB100:724 encodedonce;159latents; discard153latent prefix",
+  "i2v": {
+    "anchor": 699,
+    "decoded_slice": [
+      1,
+      25
+    ],
+    "controls": [
+      700,
+      724
+    ],
+    "right_noops": 9
+  },
+  "normalization_change": "weighted modes apply scales and accumulate RMS in fp32",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "parameter_counts": {
+    "generator_total": 1619242304,
+    "generator_trainable": 0,
+    "vae_total": 126892531,
+    "clip_total": 1193014785
+  },
+  "action_norm_instances": 60
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85091"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_full_weighted_job85088_task2/probe.py

@@ -0,0 +1,113 @@
+"""Frozen Matrix diagnostics; raw-RGB PSNR is separate from benchmark scoring."""
+import argparse
+import json
+import sys
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(ROOT / 'shared'), str(ROOT / 'Matrix-Game/Matrix-Game-2')]
+from protocol import load_minecraft_case, select_cases
+from worldmem_adapter import MatrixGameAdapter, map_actions
+from demo_utils.constant import ZERO_VAE_CACHE
+
+MODES = ['wan_oracle', 'full_original', 'full_weighted', 'i2v_original', 'i2v_weighted']
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('--mode', choices=MODES, required=True)
+parser.add_argument('--data-root', type=Path, required=True)
+parser.add_argument('--output', type=Path, required=True)
+parser.add_argument('--dry-run', action='store_true')
+
+def weighted_norm(self, x):
+    y = x.float()
+    return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight
+
+def save_result(output, name, prediction, reference):
+    prediction = prediction.float().cpu()
+    if prediction.shape != reference.shape or not torch.isfinite(prediction).all():
+        raise ValueError(f'Invalid Matrix prediction: {name}, shape={tuple(prediction.shape)}')
+    prediction = prediction.clamp(0, 1)
+    mse = (prediction - reference).square().mean((1, 2, 3))
+    report = {'case': name, 'raw_rgb_psnr': (-10 * mse.clamp_min(1e-12).log10()).tolist(),
+              'raw_rgb_mse': mse.tolist(), 'mean_rgb': prediction.mean((1, 2, 3)).tolist()}
+    (output / f'{name}.json').write_text(json.dumps(report, indent=2))
+    grid = torch.cat([torch.cat([v[t] for t in [0, 1, 2, 5, 11, 23]], dim=2)
+                      for v in (reference, prediction)], dim=1)
+    Image.fromarray(grid.mul(255).byte().permute(1, 2, 0).numpy()).save(output / f'{name}.jpg')
+    print(json.dumps({'case': name, 'mean_psnr': sum(report['raw_rgb_psnr']) / 24,
+                      'first_psnr': report['raw_rgb_psnr'][:3]}), flush=True)
+
+@torch.inference_mode()
+def main():
+    args = parser.parse_args()
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    manifest = {'diagnostic_only': True, 'mode': args.mode, 'checkpoint': str(checkpoint),
+                'weight_source': 'released raw distilled', 'optimizer': None, 'lr': None, 'lora': None,
+                'activation_checkpointing': False, 'batch_size': 1, 'trainable_parameters': 0,
+                'denoising_steps': 3, 'shift': 5, 'local_attention_latents': 6,
+                'selector': 'native sliding cache', 'reference': 'raw RGB resized352x640, not benchmark Oasis references',
+                'targets': [700, 724], 'full_history': [100, 700], 'full_prefix_latents': 153,
+                'oracle': '9copies+RGB100:724 encodedonce;159latents; discard153latent prefix',
+                'i2v': {'anchor': 699, 'decoded_slice': [1, 25], 'controls': [700, 724], 'right_noops': 9},
+                'normalization_change': 'weighted modes apply scales and accumulate RMS in fp32',
+                'cases': [{'index': r['global_test_index'], 'path': r['relative_path'],
+                           'seed': 42000000 + int(r['global_test_index'])} for r in rows]}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    from wan.modules.action_module import WanRMSNorm
+    adapter = MatrixGameAdapter(checkpoint)
+    manifest['parameter_counts'] = adapter.parameter_counts
+    manifest['action_norm_instances'] = sum(isinstance(m, WanRMSNorm) for m in adapter.pipeline.modules())
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    if args.mode.endswith('weighted'):
+        WanRMSNorm.forward = weighted_norm
+    for row in rows:
+        row.update(target_end='724', target_frames='24')
+        case = load_minecraft_case(row, args.data_root)
+        reference = F.interpolate(case['rgb'][600:624], (352, 640), mode='bilinear', align_corners=False, antialias=True)
+        seed = 42000000 + int(row['global_test_index'])
+        if args.mode == 'wan_oracle':
+            rgb = F.interpolate(case['rgb'].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            rgb = torch.cat([rgb[:1].expand(9, -1, -1, -1), rgb])
+            latent = adapter.vae.encode(rgb.transpose(0, 1)[None].mul(2).sub(1).bfloat16(), device=adapter.device, tiled=False)
+            cache, chunks = [None] * len(ZERO_VAE_CACHE), []
+            for start in range(0, latent.shape[2], 3):
+                decoded, cache = adapter.pipeline.vae_decoder(latent[:, :, start:start + 3].transpose(1, 2).half(), *cache)
+                if start >= 153:
+                    chunks.append(decoded.cpu())
+            prediction = torch.cat(chunks, dim=1)[0].float().add(1).div(2)
+            del latent, rgb, cache, decoded, chunks
+        elif args.mode.startswith('full'):
+            prediction, _ = adapter.generate(case, target_frames=24, seed=seed)
+        else:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            anchor = F.interpolate(case['rgb'][599:600].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            anchor = anchor.transpose(0, 1)[None].mul(2).sub(1).bfloat16()
+            image = anchor.new_zeros((1, 3, 33, 352, 640))
+            image[:, :, :1] = anchor
+            image_latents = adapter.vae.encode(image, device=adapter.device, tiled=False)
+            mask = torch.zeros_like(image_latents[:, :4])
+            mask[:, :, :1] = 1
+            actions = torch.cat([case['actions'][600:624], torch.zeros(9, 25)]).cuda().bfloat16()
+            keyboard, mouse = map_actions(actions, 1.0)
+            condition = {'cond_concat': torch.cat([mask, image_latents], dim=1),
+                         'visual_context': adapter.vae.clip.encode_video(anchor).bfloat16(),
+                         'keyboard_cond': keyboard[None], 'mouse_cond': mouse[None]}
+            noise = torch.randn((1, 16, 9, 44, 80), device='cuda', dtype=torch.bfloat16)
+            videos = adapter.pipeline.inference(noise, condition, initial_latent=None, mode='universal')
+            prediction = torch.cat(videos, dim=1)[0, 1:25].float().add(1).div(2)
+            del image, anchor, image_latents, condition, videos, noise
+        save_result(args.output, f"case{row['global_test_index']}_{args.mode}", prediction, reference)
+        del prediction, reference, case
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'cases': 2, 'frames': 24}))
+    print('DIAGNOSTICS_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/case0_i2v_original.json

@@ -0,0 +1,81 @@
+{
+  "case": "case0_i2v_original",
+  "raw_rgb_psnr": [
+    28.51933479309082,
+    5.62062931060791,
+    6.261198997497559,
+    7.325035095214844,
+    7.525331497192383,
+    7.780514240264893,
+    7.777151107788086,
+    7.821002006530762,
+    7.603299617767334,
+    7.544827461242676,
+    7.494100570678711,
+    7.490266799926758,
+    7.607268810272217,
+    7.717894554138184,
+    7.877525329589844,
+    7.909063816070557,
+    7.8667497634887695,
+    7.888334274291992,
+    7.912131309509277,
+    7.934691905975342,
+    7.701882839202881,
+    6.664777755737305,
+    7.160373687744141,
+    7.919079780578613
+  ],
+  "raw_rgb_mse": [
+    0.001406263094395399,
+    0.27411770820617676,
+    0.23652665317058563,
+    0.18513838946819305,
+    0.17679372429847717,
+    0.16670498251914978,
+    0.16683414578437805,
+    0.16515807807445526,
+    0.17364810407161713,
+    0.17600186169147491,
+    0.17806968092918396,
+    0.17822694778442383,
+    0.17348946630954742,
+    0.16912606358528137,
+    0.1630224734544754,
+    0.16184288263320923,
+    0.1634274572134018,
+    0.16261722147464752,
+    0.1617286056280136,
+    0.16089065372943878,
+    0.16975075006484985,
+    0.2155371904373169,
+    0.19229261577129364,
+    0.16147007048130035
+  ],
+  "mean_rgb": [
+    0.7052807807922363,
+    0.22267355024814606,
+    0.26383358240127563,
+    0.31955206394195557,
+    0.3294263184070587,
+    0.3413337469100952,
+    0.34034264087677,
+    0.3422088921070099,
+    0.3322543799877167,
+    0.3306882083415985,
+    0.32795271277427673,
+    0.32645344734191895,
+    0.32989341020584106,
+    0.33399924635887146,
+    0.34038686752319336,
+    0.34092986583709717,
+    0.33980700373649597,
+    0.3406497836112976,
+    0.3414357602596283,
+    0.3423357307910919,
+    0.3307926654815674,
+    0.2806178331375122,
+    0.30535194277763367,
+    0.34129393100738525
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/case1_i2v_original.json

@@ -0,0 +1,81 @@
+{
+  "case": "case1_i2v_original",
+  "raw_rgb_psnr": [
+    15.054824829101562,
+    11.44872760772705,
+    9.014760971069336,
+    7.8256916999816895,
+    6.933335781097412,
+    6.524610996246338,
+    6.566136360168457,
+    6.591427803039551,
+    6.48160982131958,
+    6.368434906005859,
+    4.373075008392334,
+    5.369548320770264,
+    5.37614631652832,
+    5.41501522064209,
+    5.501518726348877,
+    5.463625907897949,
+    5.47122859954834,
+    5.503015041351318,
+    5.575506210327148,
+    5.547548294067383,
+    5.560653209686279,
+    5.55626106262207,
+    5.5676751136779785,
+    5.580542087554932
+  ],
+  "raw_rgb_mse": [
+    0.031226083636283875,
+    0.07163531333208084,
+    0.12546537816524506,
+    0.16497983038425446,
+    0.20261257886886597,
+    0.22260703146457672,
+    0.22048871219158173,
+    0.21920840442180634,
+    0.22482208907604218,
+    0.23075786232948303,
+    0.36533603072166443,
+    0.29043248295783997,
+    0.289991557598114,
+    0.2874077558517456,
+    0.281739741563797,
+    0.284208744764328,
+    0.28371164202690125,
+    0.28164270520210266,
+    0.27698060870170593,
+    0.27876943349838257,
+    0.27792951464653015,
+    0.27821075916290283,
+    0.27748051285743713,
+    0.27665963768959045
+  ],
+  "mean_rgb": [
+    0.6844472289085388,
+    0.6284617781639099,
+    0.5540845394134521,
+    0.49505874514579773,
+    0.4436206817626953,
+    0.41198185086250305,
+    0.4130478501319885,
+    0.41542676091194153,
+    0.406199187040329,
+    0.39254048466682434,
+    0.16926062107086182,
+    0.22599884867668152,
+    0.22546076774597168,
+    0.22791211307048798,
+    0.23338943719863892,
+    0.23143425583839417,
+    0.23280683159828186,
+    0.23524777591228485,
+    0.23973938822746277,
+    0.23811960220336914,
+    0.23764897882938385,
+    0.2369965761899948,
+    0.23754198849201202,
+    0.23855453729629517
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/complete.json

@@ -0,0 +1 @@
+{"mode": "i2v_original", "cases": 2, "frames": 24}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/launcher.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-diagnostic
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=01:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(wan_oracle full_original full_weighted i2v_original i2v_weighted)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/diagnose_matrix_game2.py" "$OUTPUT/probe.py"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/diagnose_matrix_game2.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/manifest.json

@@ -0,0 +1,59 @@
+{
+  "diagnostic_only": true,
+  "mode": "i2v_original",
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "weight_source": "released raw distilled",
+  "optimizer": null,
+  "lr": null,
+  "lora": null,
+  "activation_checkpointing": false,
+  "batch_size": 1,
+  "trainable_parameters": 0,
+  "denoising_steps": 3,
+  "shift": 5,
+  "local_attention_latents": 6,
+  "selector": "native sliding cache",
+  "reference": "raw RGB resized352x640, not benchmark Oasis references",
+  "targets": [
+    700,
+    724
+  ],
+  "full_history": [
+    100,
+    700
+  ],
+  "full_prefix_latents": 153,
+  "oracle": "9copies+RGB100:724 encodedonce;159latents; discard153latent prefix",
+  "i2v": {
+    "anchor": 699,
+    "decoded_slice": [
+      1,
+      25
+    ],
+    "controls": [
+      700,
+      724
+    ],
+    "right_noops": 9
+  },
+  "normalization_change": "weighted modes apply scales and accumulate RMS in fp32",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "parameter_counts": {
+    "generator_total": 1619242304,
+    "generator_trainable": 0,
+    "vae_total": 126892531,
+    "clip_total": 1193014785
+  },
+  "action_norm_instances": 60
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85092"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_original_job85088_task3/probe.py

@@ -0,0 +1,113 @@
+"""Frozen Matrix diagnostics; raw-RGB PSNR is separate from benchmark scoring."""
+import argparse
+import json
+import sys
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(ROOT / 'shared'), str(ROOT / 'Matrix-Game/Matrix-Game-2')]
+from protocol import load_minecraft_case, select_cases
+from worldmem_adapter import MatrixGameAdapter, map_actions
+from demo_utils.constant import ZERO_VAE_CACHE
+
+MODES = ['wan_oracle', 'full_original', 'full_weighted', 'i2v_original', 'i2v_weighted']
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('--mode', choices=MODES, required=True)
+parser.add_argument('--data-root', type=Path, required=True)
+parser.add_argument('--output', type=Path, required=True)
+parser.add_argument('--dry-run', action='store_true')
+
+def weighted_norm(self, x):
+    y = x.float()
+    return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight
+
+def save_result(output, name, prediction, reference):
+    prediction = prediction.float().cpu()
+    if prediction.shape != reference.shape or not torch.isfinite(prediction).all():
+        raise ValueError(f'Invalid Matrix prediction: {name}, shape={tuple(prediction.shape)}')
+    prediction = prediction.clamp(0, 1)
+    mse = (prediction - reference).square().mean((1, 2, 3))
+    report = {'case': name, 'raw_rgb_psnr': (-10 * mse.clamp_min(1e-12).log10()).tolist(),
+              'raw_rgb_mse': mse.tolist(), 'mean_rgb': prediction.mean((1, 2, 3)).tolist()}
+    (output / f'{name}.json').write_text(json.dumps(report, indent=2))
+    grid = torch.cat([torch.cat([v[t] for t in [0, 1, 2, 5, 11, 23]], dim=2)
+                      for v in (reference, prediction)], dim=1)
+    Image.fromarray(grid.mul(255).byte().permute(1, 2, 0).numpy()).save(output / f'{name}.jpg')
+    print(json.dumps({'case': name, 'mean_psnr': sum(report['raw_rgb_psnr']) / 24,
+                      'first_psnr': report['raw_rgb_psnr'][:3]}), flush=True)
+
+@torch.inference_mode()
+def main():
+    args = parser.parse_args()
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    manifest = {'diagnostic_only': True, 'mode': args.mode, 'checkpoint': str(checkpoint),
+                'weight_source': 'released raw distilled', 'optimizer': None, 'lr': None, 'lora': None,
+                'activation_checkpointing': False, 'batch_size': 1, 'trainable_parameters': 0,
+                'denoising_steps': 3, 'shift': 5, 'local_attention_latents': 6,
+                'selector': 'native sliding cache', 'reference': 'raw RGB resized352x640, not benchmark Oasis references',
+                'targets': [700, 724], 'full_history': [100, 700], 'full_prefix_latents': 153,
+                'oracle': '9copies+RGB100:724 encodedonce;159latents; discard153latent prefix',
+                'i2v': {'anchor': 699, 'decoded_slice': [1, 25], 'controls': [700, 724], 'right_noops': 9},
+                'normalization_change': 'weighted modes apply scales and accumulate RMS in fp32',
+                'cases': [{'index': r['global_test_index'], 'path': r['relative_path'],
+                           'seed': 42000000 + int(r['global_test_index'])} for r in rows]}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    from wan.modules.action_module import WanRMSNorm
+    adapter = MatrixGameAdapter(checkpoint)
+    manifest['parameter_counts'] = adapter.parameter_counts
+    manifest['action_norm_instances'] = sum(isinstance(m, WanRMSNorm) for m in adapter.pipeline.modules())
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    if args.mode.endswith('weighted'):
+        WanRMSNorm.forward = weighted_norm
+    for row in rows:
+        row.update(target_end='724', target_frames='24')
+        case = load_minecraft_case(row, args.data_root)
+        reference = F.interpolate(case['rgb'][600:624], (352, 640), mode='bilinear', align_corners=False, antialias=True)
+        seed = 42000000 + int(row['global_test_index'])
+        if args.mode == 'wan_oracle':
+            rgb = F.interpolate(case['rgb'].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            rgb = torch.cat([rgb[:1].expand(9, -1, -1, -1), rgb])
+            latent = adapter.vae.encode(rgb.transpose(0, 1)[None].mul(2).sub(1).bfloat16(), device=adapter.device, tiled=False)
+            cache, chunks = [None] * len(ZERO_VAE_CACHE), []
+            for start in range(0, latent.shape[2], 3):
+                decoded, cache = adapter.pipeline.vae_decoder(latent[:, :, start:start + 3].transpose(1, 2).half(), *cache)
+                if start >= 153:
+                    chunks.append(decoded.cpu())
+            prediction = torch.cat(chunks, dim=1)[0].float().add(1).div(2)
+            del latent, rgb, cache, decoded, chunks
+        elif args.mode.startswith('full'):
+            prediction, _ = adapter.generate(case, target_frames=24, seed=seed)
+        else:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            anchor = F.interpolate(case['rgb'][599:600].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            anchor = anchor.transpose(0, 1)[None].mul(2).sub(1).bfloat16()
+            image = anchor.new_zeros((1, 3, 33, 352, 640))
+            image[:, :, :1] = anchor
+            image_latents = adapter.vae.encode(image, device=adapter.device, tiled=False)
+            mask = torch.zeros_like(image_latents[:, :4])
+            mask[:, :, :1] = 1
+            actions = torch.cat([case['actions'][600:624], torch.zeros(9, 25)]).cuda().bfloat16()
+            keyboard, mouse = map_actions(actions, 1.0)
+            condition = {'cond_concat': torch.cat([mask, image_latents], dim=1),
+                         'visual_context': adapter.vae.clip.encode_video(anchor).bfloat16(),
+                         'keyboard_cond': keyboard[None], 'mouse_cond': mouse[None]}
+            noise = torch.randn((1, 16, 9, 44, 80), device='cuda', dtype=torch.bfloat16)
+            videos = adapter.pipeline.inference(noise, condition, initial_latent=None, mode='universal')
+            prediction = torch.cat(videos, dim=1)[0, 1:25].float().add(1).div(2)
+            del image, anchor, image_latents, condition, videos, noise
+        save_result(args.output, f"case{row['global_test_index']}_{args.mode}", prediction, reference)
+        del prediction, reference, case
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'cases': 2, 'frames': 24}))
+    print('DIAGNOSTICS_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/case0_i2v_weighted.json

@@ -0,0 +1,81 @@
+{
+  "case": "case0_i2v_weighted",
+  "raw_rgb_psnr": [
+    28.400230407714844,
+    5.406888008117676,
+    6.059108257293701,
+    7.271352767944336,
+    7.498164653778076,
+    7.74724006652832,
+    7.7596211433410645,
+    7.7729315757751465,
+    7.52515983581543,
+    7.486701488494873,
+    7.524697303771973,
+    7.638741493225098,
+    7.781850814819336,
+    7.818033695220947,
+    7.900123119354248,
+    7.916782379150391,
+    7.866796493530273,
+    7.875070095062256,
+    7.894111633300781,
+    7.932208061218262,
+    7.702528953552246,
+    7.582130432128906,
+    7.766695976257324,
+    7.993913650512695
+  ],
+  "raw_rgb_mse": [
+    0.0014453631592914462,
+    0.28794607520103455,
+    0.2477930635213852,
+    0.18744103610515594,
+    0.17790310084819794,
+    0.16798712313175201,
+    0.16750890016555786,
+    0.1669963002204895,
+    0.17680072784423828,
+    0.1783733069896698,
+    0.1768195480108261,
+    0.17223675549030304,
+    0.16665369272232056,
+    0.16527098417282104,
+    0.16217641532421112,
+    0.16155549883842468,
+    0.16342569887638092,
+    0.16311465203762054,
+    0.16240103542804718,
+    0.16098269820213318,
+    0.16972550749778748,
+    0.1744966059923172,
+    0.16723622381687164,
+    0.15871158242225647
+  ],
+  "mean_rgb": [
+    0.7053292989730835,
+    0.20904700458049774,
+    0.2530539929866791,
+    0.3162554204463959,
+    0.3277013599872589,
+    0.33897286653518677,
+    0.3385704457759857,
+    0.3386596143245697,
+    0.32705172896385193,
+    0.3246579170227051,
+    0.32665011286735535,
+    0.33084917068481445,
+    0.33509644865989685,
+    0.335825651884079,
+    0.34004831314086914,
+    0.34061479568481445,
+    0.338895708322525,
+    0.339072585105896,
+    0.33963221311569214,
+    0.3410741984844208,
+    0.3300783038139343,
+    0.32414621114730835,
+    0.33280378580093384,
+    0.3427526652812958
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/case1_i2v_weighted.json

@@ -0,0 +1,81 @@
+{
+  "case": "case1_i2v_weighted",
+  "raw_rgb_psnr": [
+    15.572078704833984,
+    11.67737865447998,
+    9.182900428771973,
+    7.9456071853637695,
+    6.916970729827881,
+    6.555791854858398,
+    6.522993564605713,
+    6.507280349731445,
+    6.314092636108398,
+    6.2913408279418945,
+    6.392267227172852,
+    6.423827171325684,
+    6.575774669647217,
+    4.436397075653076,
+    4.66413688659668,
+    4.843669414520264,
+    4.8888068199157715,
+    4.991722106933594,
+    5.058438301086426,
+    5.032224178314209,
+    5.076114654541016,
+    5.092604637145996,
+    5.115626335144043,
+    5.137517929077148
+  ],
+  "raw_rgb_mse": [
+    0.027719931676983833,
+    0.06796137243509293,
+    0.12070075422525406,
+    0.16048678755760193,
+    0.20337752997875214,
+    0.22101451456546783,
+    0.22268995642662048,
+    0.22349712252616882,
+    0.2336634248495102,
+    0.23489077389240265,
+    0.22949503362178802,
+    0.2278333306312561,
+    0.21999992430210114,
+    0.3600479066371918,
+    0.34165385365486145,
+    0.32781821489334106,
+    0.3244287371635437,
+    0.3168310821056366,
+    0.31200113892555237,
+    0.31389009952545166,
+    0.3107338547706604,
+    0.3095562160015106,
+    0.30791962146759033,
+    0.3063713610172272
+  ],
+  "mean_rgb": [
+    0.6906836032867432,
+    0.6348850727081299,
+    0.5620357394218445,
+    0.5025600790977478,
+    0.4471380114555359,
+    0.41820263862609863,
+    0.4126875400543213,
+    0.41116610169410706,
+    0.3962208926677704,
+    0.3935016095638275,
+    0.3975541591644287,
+    0.39638906717300415,
+    0.40340858697891235,
+    0.16516827046871185,
+    0.18019184470176697,
+    0.19274631142616272,
+    0.1964321732521057,
+    0.20392490923404694,
+    0.20804454386234283,
+    0.2059834748506546,
+    0.20835448801517487,
+    0.20949353277683258,
+    0.21102812886238098,
+    0.2126505970954895
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/complete.json

@@ -0,0 +1 @@
+{"mode": "i2v_weighted", "cases": 2, "frames": 24}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/launcher.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-diagnostic
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=01:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(wan_oracle full_original full_weighted i2v_original i2v_weighted)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/diagnose_matrix_game2.py" "$OUTPUT/probe.py"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/diagnose_matrix_game2.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/manifest.json

@@ -0,0 +1,59 @@
+{
+  "diagnostic_only": true,
+  "mode": "i2v_weighted",
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "weight_source": "released raw distilled",
+  "optimizer": null,
+  "lr": null,
+  "lora": null,
+  "activation_checkpointing": false,
+  "batch_size": 1,
+  "trainable_parameters": 0,
+  "denoising_steps": 3,
+  "shift": 5,
+  "local_attention_latents": 6,
+  "selector": "native sliding cache",
+  "reference": "raw RGB resized352x640, not benchmark Oasis references",
+  "targets": [
+    700,
+    724
+  ],
+  "full_history": [
+    100,
+    700
+  ],
+  "full_prefix_latents": 153,
+  "oracle": "9copies+RGB100:724 encodedonce;159latents; discard153latent prefix",
+  "i2v": {
+    "anchor": 699,
+    "decoded_slice": [
+      1,
+      25
+    ],
+    "controls": [
+      700,
+      724
+    ],
+    "right_noops": 9
+  },
+  "normalization_change": "weighted modes apply scales and accumulate RMS in fp32",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "parameter_counts": {
+    "generator_total": 1619242304,
+    "generator_trainable": 0,
+    "vae_total": 126892531,
+    "clip_total": 1193014785
+  },
+  "action_norm_instances": 60
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85088"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_i2v_weighted_job85088_task4/probe.py

@@ -0,0 +1,113 @@
+"""Frozen Matrix diagnostics; raw-RGB PSNR is separate from benchmark scoring."""
+import argparse
+import json
+import sys
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(ROOT / 'shared'), str(ROOT / 'Matrix-Game/Matrix-Game-2')]
+from protocol import load_minecraft_case, select_cases
+from worldmem_adapter import MatrixGameAdapter, map_actions
+from demo_utils.constant import ZERO_VAE_CACHE
+
+MODES = ['wan_oracle', 'full_original', 'full_weighted', 'i2v_original', 'i2v_weighted']
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('--mode', choices=MODES, required=True)
+parser.add_argument('--data-root', type=Path, required=True)
+parser.add_argument('--output', type=Path, required=True)
+parser.add_argument('--dry-run', action='store_true')
+
+def weighted_norm(self, x):
+    y = x.float()
+    return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight
+
+def save_result(output, name, prediction, reference):
+    prediction = prediction.float().cpu()
+    if prediction.shape != reference.shape or not torch.isfinite(prediction).all():
+        raise ValueError(f'Invalid Matrix prediction: {name}, shape={tuple(prediction.shape)}')
+    prediction = prediction.clamp(0, 1)
+    mse = (prediction - reference).square().mean((1, 2, 3))
+    report = {'case': name, 'raw_rgb_psnr': (-10 * mse.clamp_min(1e-12).log10()).tolist(),
+              'raw_rgb_mse': mse.tolist(), 'mean_rgb': prediction.mean((1, 2, 3)).tolist()}
+    (output / f'{name}.json').write_text(json.dumps(report, indent=2))
+    grid = torch.cat([torch.cat([v[t] for t in [0, 1, 2, 5, 11, 23]], dim=2)
+                      for v in (reference, prediction)], dim=1)
+    Image.fromarray(grid.mul(255).byte().permute(1, 2, 0).numpy()).save(output / f'{name}.jpg')
+    print(json.dumps({'case': name, 'mean_psnr': sum(report['raw_rgb_psnr']) / 24,
+                      'first_psnr': report['raw_rgb_psnr'][:3]}), flush=True)
+
+@torch.inference_mode()
+def main():
+    args = parser.parse_args()
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    manifest = {'diagnostic_only': True, 'mode': args.mode, 'checkpoint': str(checkpoint),
+                'weight_source': 'released raw distilled', 'optimizer': None, 'lr': None, 'lora': None,
+                'activation_checkpointing': False, 'batch_size': 1, 'trainable_parameters': 0,
+                'denoising_steps': 3, 'shift': 5, 'local_attention_latents': 6,
+                'selector': 'native sliding cache', 'reference': 'raw RGB resized352x640, not benchmark Oasis references',
+                'targets': [700, 724], 'full_history': [100, 700], 'full_prefix_latents': 153,
+                'oracle': '9copies+RGB100:724 encodedonce;159latents; discard153latent prefix',
+                'i2v': {'anchor': 699, 'decoded_slice': [1, 25], 'controls': [700, 724], 'right_noops': 9},
+                'normalization_change': 'weighted modes apply scales and accumulate RMS in fp32',
+                'cases': [{'index': r['global_test_index'], 'path': r['relative_path'],
+                           'seed': 42000000 + int(r['global_test_index'])} for r in rows]}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    from wan.modules.action_module import WanRMSNorm
+    adapter = MatrixGameAdapter(checkpoint)
+    manifest['parameter_counts'] = adapter.parameter_counts
+    manifest['action_norm_instances'] = sum(isinstance(m, WanRMSNorm) for m in adapter.pipeline.modules())
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    if args.mode.endswith('weighted'):
+        WanRMSNorm.forward = weighted_norm
+    for row in rows:
+        row.update(target_end='724', target_frames='24')
+        case = load_minecraft_case(row, args.data_root)
+        reference = F.interpolate(case['rgb'][600:624], (352, 640), mode='bilinear', align_corners=False, antialias=True)
+        seed = 42000000 + int(row['global_test_index'])
+        if args.mode == 'wan_oracle':
+            rgb = F.interpolate(case['rgb'].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            rgb = torch.cat([rgb[:1].expand(9, -1, -1, -1), rgb])
+            latent = adapter.vae.encode(rgb.transpose(0, 1)[None].mul(2).sub(1).bfloat16(), device=adapter.device, tiled=False)
+            cache, chunks = [None] * len(ZERO_VAE_CACHE), []
+            for start in range(0, latent.shape[2], 3):
+                decoded, cache = adapter.pipeline.vae_decoder(latent[:, :, start:start + 3].transpose(1, 2).half(), *cache)
+                if start >= 153:
+                    chunks.append(decoded.cpu())
+            prediction = torch.cat(chunks, dim=1)[0].float().add(1).div(2)
+            del latent, rgb, cache, decoded, chunks
+        elif args.mode.startswith('full'):
+            prediction, _ = adapter.generate(case, target_frames=24, seed=seed)
+        else:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            anchor = F.interpolate(case['rgb'][599:600].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True)
+            anchor = anchor.transpose(0, 1)[None].mul(2).sub(1).bfloat16()
+            image = anchor.new_zeros((1, 3, 33, 352, 640))
+            image[:, :, :1] = anchor
+            image_latents = adapter.vae.encode(image, device=adapter.device, tiled=False)
+            mask = torch.zeros_like(image_latents[:, :4])
+            mask[:, :, :1] = 1
+            actions = torch.cat([case['actions'][600:624], torch.zeros(9, 25)]).cuda().bfloat16()
+            keyboard, mouse = map_actions(actions, 1.0)
+            condition = {'cond_concat': torch.cat([mask, image_latents], dim=1),
+                         'visual_context': adapter.vae.clip.encode_video(anchor).bfloat16(),
+                         'keyboard_cond': keyboard[None], 'mouse_cond': mouse[None]}
+            noise = torch.randn((1, 16, 9, 44, 80), device='cuda', dtype=torch.bfloat16)
+            videos = adapter.pipeline.inference(noise, condition, initial_latent=None, mode='universal')
+            prediction = torch.cat(videos, dim=1)[0, 1:25].float().add(1).div(2)
+            del image, anchor, image_latents, condition, videos, noise
+        save_result(args.output, f"case{row['global_test_index']}_{args.mode}", prediction, reference)
+        del prediction, reference, case
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'cases': 2, 'frames': 24}))
+    print('DIAGNOSTICS_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/case0/comparison.json

@@ -0,0 +1,81 @@
+{
+  "case": "comparison",
+  "raw_rgb_psnr": [
+    28.761943817138672,
+    24.894990921020508,
+    5.066094875335693,
+    6.884883880615234,
+    7.142206192016602,
+    7.228761196136475,
+    7.280294895172119,
+    7.196128845214844,
+    7.207036972045898,
+    7.207610130310059,
+    7.2647624015808105,
+    7.177729606628418,
+    7.162692546844482,
+    7.2452826499938965,
+    7.280104637145996,
+    7.2675628662109375,
+    7.268505096435547,
+    7.322355270385742,
+    7.28434944152832,
+    7.270142555236816,
+    7.088680267333984,
+    6.997251510620117,
+    6.956180572509766,
+    7.0359296798706055
+  ],
+  "raw_rgb_mse": [
+    0.0013298585545271635,
+    0.0032396710012108088,
+    0.3114515542984009,
+    0.20488570630550385,
+    0.19309870898723602,
+    0.18928833305835724,
+    0.18705551326274872,
+    0.1907159984111786,
+    0.19023756682872772,
+    0.19021247327327728,
+    0.1877257078886032,
+    0.1915256828069687,
+    0.1921899914741516,
+    0.18856962025165558,
+    0.18706370890140533,
+    0.18760469555854797,
+    0.18756400048732758,
+    0.18525265157222748,
+    0.186880961060524,
+    0.18749330937862396,
+    0.19549334049224854,
+    0.1996525377035141,
+    0.20154961943626404,
+    0.19788235425949097
+  ],
+  "mean_rgb": [
+    0.7190113663673401,
+    0.7229689359664917,
+    0.19768567383289337,
+    0.296479195356369,
+    0.30714768171310425,
+    0.31132906675338745,
+    0.31336742639541626,
+    0.3084028661251068,
+    0.3105626106262207,
+    0.3106946647167206,
+    0.3127269446849823,
+    0.3085455596446991,
+    0.30667993426322937,
+    0.3099704384803772,
+    0.31143423914909363,
+    0.31037086248397827,
+    0.31108224391937256,
+    0.3149312436580658,
+    0.3144523501396179,
+    0.31461331248283386,
+    0.30496296286582947,
+    0.3001496493816376,
+    0.2973284423351288,
+    0.30110853910446167
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/case0/inputs.npz

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

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/case1/comparison.json

@@ -0,0 +1,81 @@
+{
+  "case": "comparison",
+  "raw_rgb_psnr": [
+    19.81435775756836,
+    18.266035079956055,
+    17.720809936523438,
+    17.051300048828125,
+    17.8521785736084,
+    17.96320152282715,
+    17.770442962646484,
+    17.66672706604004,
+    17.461271286010742,
+    17.219820022583008,
+    17.17270851135254,
+    17.01397132873535,
+    16.4091739654541,
+    16.703092575073242,
+    17.036758422851562,
+    16.95575523376465,
+    16.949073791503906,
+    16.948787689208984,
+    16.964191436767578,
+    17.009931564331055,
+    16.958953857421875,
+    16.964096069335938,
+    16.84088134765625,
+    16.837522506713867
+  ],
+  "raw_rgb_mse": [
+    0.010436725802719593,
+    0.014907212927937508,
+    0.01690126210451126,
+    0.019718321040272713,
+    0.016397668048739433,
+    0.01598379760980606,
+    0.016709202900528908,
+    0.01711304485797882,
+    0.01794208399951458,
+    0.018967844545841217,
+    0.01917472667992115,
+    0.019888538867235184,
+    0.02286033146083355,
+    0.021364402025938034,
+    0.019784456118941307,
+    0.02015693485736847,
+    0.020187973976135254,
+    0.020189296454191208,
+    0.02011781372129917,
+    0.019907046109437943,
+    0.020142091438174248,
+    0.020118258893489838,
+    0.0206972137093544,
+    0.02071322314441204
+  ],
+  "mean_rgb": [
+    0.7137359976768494,
+    0.7148902416229248,
+    0.7157489061355591,
+    0.7161552906036377,
+    0.713788628578186,
+    0.7127888202667236,
+    0.7122524976730347,
+    0.7112450003623962,
+    0.7104371786117554,
+    0.7100856900215149,
+    0.7079524397850037,
+    0.7049611210823059,
+    0.700785756111145,
+    0.7041963934898376,
+    0.7078825831413269,
+    0.7077915668487549,
+    0.7052977085113525,
+    0.7040800452232361,
+    0.7037867903709412,
+    0.7034406065940857,
+    0.7064051628112793,
+    0.7068390846252441,
+    0.7078206539154053,
+    0.7080984115600586
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/case1/inputs.npz

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

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/complete.json

@@ -0,0 +1 @@
+{"mode": "demo_actions", "complete": true}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/inference.py

@@ -0,0 +1,169 @@
+import os
+import argparse
+import torch
+import numpy as np
+
+from omegaconf import OmegaConf
+from torchvision.transforms import v2
+from diffusers.utils import load_image
+from einops import rearrange
+from pipeline import CausalInferencePipeline
+from wan.vae.wanx_vae import get_wanx_vae_wrapper
+from demo_utils.vae_block3 import VAEDecoderWrapper
+from utils.visualize import process_video
+from utils.misc import set_seed
+from utils.conditions import *
+from utils.wan_wrapper import WanDiffusionWrapper
+from safetensors.torch import load_file
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, default="configs/inference_yaml/inference_universal.yaml", help="Path to the config file")
+    parser.add_argument("--checkpoint_path", type=str, default="", help="Path to the checkpoint")
+    parser.add_argument("--img_path", type=str, default="demo_images/universal/0000.png", help="Path to the image")
+    parser.add_argument("--output_folder", type=str, default="outputs/", help="Output folder")
+    parser.add_argument("--num_output_frames", type=int, default=150,
+                        help="Number of output latent frames")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--pretrained_model_path", type=str, default="Matrix-Game-2.0", help="Path to the VAE model folder")
+    args = parser.parse_args()
+    return args
+
+class InteractiveGameInference:
+    def __init__(self, args):
+        self.args = args
+        self.device = torch.device("cuda")
+        self.weight_dtype = torch.bfloat16
+
+        self._init_config()
+        self._init_models()
+
+        self.frame_process = v2.Compose([
+            v2.Resize(size=(352, 640), antialias=True),
+            v2.ToTensor(),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+
+    def _init_config(self):
+        self.config = OmegaConf.load(self.args.config_path)
+
+    def _init_models(self):
+        # Initialize pipeline
+        generator = WanDiffusionWrapper(
+            **getattr(self.config, "model_kwargs", {}), is_causal=True)
+        current_vae_decoder = VAEDecoderWrapper()
+        vae_state_dict = torch.load(os.path.join(self.args.pretrained_model_path, "Wan2.1_VAE.pth"), map_location="cpu")
+        decoder_state_dict = {}
+        for key, value in vae_state_dict.items():
+            if 'decoder.' in key or 'conv2' in key:
+                decoder_state_dict[key] = value
+        current_vae_decoder.load_state_dict(decoder_state_dict)
+        current_vae_decoder.to(self.device, torch.float16)
+        current_vae_decoder.requires_grad_(False)
+        current_vae_decoder.eval()
+        current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")
+        pipeline = CausalInferencePipeline(self.config, generator=generator, vae_decoder=current_vae_decoder)
+        if self.args.checkpoint_path:
+            print("Loading Pretrained Model...")
+            state_dict = load_file(self.args.checkpoint_path)
+            pipeline.generator.load_state_dict(state_dict)
+
+        self.pipeline = pipeline.to(device=self.device, dtype=self.weight_dtype)
+        self.pipeline.vae_decoder.to(torch.float16)
+
+        vae = get_wanx_vae_wrapper(self.args.pretrained_model_path, torch.float16)
+        vae.requires_grad_(False)
+        vae.eval()
+        self.vae = vae.to(self.device, self.weight_dtype)
+
+    def _resizecrop(self, image, th, tw):
+        w, h = image.size
+        if h / w > th / tw:
+            new_w = int(w)
+            new_h = int(new_w * th / tw)
+        else:
+            new_h = int(h)
+            new_w = int(new_h * tw / th)
+        left = (w - new_w) / 2
+        top = (h - new_h) / 2
+        right = (w + new_w) / 2
+        bottom = (h + new_h) / 2
+        image = image.crop((left, top, right, bottom))
+        return image
+    
+    def generate_videos(self):
+        mode = self.config.pop('mode')
+        assert mode in ['universal', 'gta_drive', 'templerun']
+
+        image = load_image(self.args.img_path)
+        image = self._resizecrop(image, 352, 640)
+        image = self.frame_process(image)[None, :, None, :, :].to(dtype=self.weight_dtype, device=self.device)
+        # Encode the input image as the first latent
+        padding_video = torch.zeros_like(image).repeat(1, 1, 4 * (self.args.num_output_frames - 1), 1, 1)
+        img_cond = torch.concat([image, padding_video], dim=2)
+        tiler_kwargs={"tiled": True, "tile_size": [44, 80], "tile_stride": [23, 38]}
+        img_cond = self.vae.encode(img_cond, device=self.device, **tiler_kwargs).to(self.device)
+        mask_cond = torch.ones_like(img_cond)
+        mask_cond[:, :, 1:] = 0
+        cond_concat = torch.cat([mask_cond[:, :4], img_cond], dim=1) 
+        visual_context = self.vae.clip.encode_video(image)
+        sampled_noise = torch.randn(
+            [1, 16,self.args.num_output_frames, 44, 80], device=self.device, dtype=self.weight_dtype
+        )
+        num_frames = (self.args.num_output_frames - 1) * 4 + 1
+        
+        conditional_dict = {
+            "cond_concat": cond_concat.to(device=self.device, dtype=self.weight_dtype),
+            "visual_context": visual_context.to(device=self.device, dtype=self.weight_dtype)
+        }
+        
+        if mode == 'universal':
+            cond_data = Bench_actions_universal(num_frames)
+            mouse_condition = cond_data['mouse_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+            conditional_dict['mouse_cond'] = mouse_condition
+        elif mode == 'gta_drive':
+            cond_data = Bench_actions_gta_drive(num_frames)
+            mouse_condition = cond_data['mouse_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+            conditional_dict['mouse_cond'] = mouse_condition
+        else:
+            cond_data = Bench_actions_templerun(num_frames)
+        keyboard_condition = cond_data['keyboard_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+        conditional_dict['keyboard_cond'] = keyboard_condition
+        
+        with torch.no_grad():
+            videos = self.pipeline.inference(
+                noise=sampled_noise,
+                conditional_dict=conditional_dict,
+                return_latents=False,
+                mode=mode,
+                profile=False
+            )
+
+        videos_tensor = torch.cat(videos, dim=1)
+        videos = rearrange(videos_tensor, "B T C H W -> B T H W C")
+        videos = ((videos.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8)[0]
+        video = np.ascontiguousarray(videos)
+        mouse_icon = 'assets/images/mouse.png'
+        if mode != 'templerun':
+            config = (
+                keyboard_condition[0].float().cpu().numpy(),
+                mouse_condition[0].float().cpu().numpy()
+            )
+        else:
+            config = (
+                keyboard_condition[0].float().cpu().numpy()
+            )
+        process_video(video.astype(np.uint8), self.args.output_folder+f'/demo.mp4', config, mouse_icon, mouse_scale=0.1, process_icon=False, mode=mode)
+        process_video(video.astype(np.uint8), self.args.output_folder+f'/demo_icon.mp4', config, mouse_icon, mouse_scale=0.1, process_icon=True, mode=mode)
+        print("Done")
+
+def main():
+    """Main entry point for video generation."""
+    args = parse_args()
+    set_seed(args.seed)
+    os.makedirs(args.output_folder, exist_ok=True)
+    pipeline = InteractiveGameInference(args)
+    pipeline.generate_videos()
+
+if __name__ == "__main__":
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/inference_universal.yaml

@@ -0,0 +1,21 @@
+denoising_step_list:
+- 1000
+- 666
+- 333
+warp_denoising_step: true
+ts_schedule: false
+mixed_precision: true
+seed: 42
+image_or_video_shape:
+- 1
+- 16
+- 15
+- 44
+- 80
+num_frame_per_block: 3
+context_noise: 0
+mode: universal
+causal: true
+model_kwargs:
+  timestep_shift: 5.0
+  model_config: configs/distilled_model/universal

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/launcher.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-native-control
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=02:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(official_demo demo_actions neutral recorded small_mouse)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_native_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export TORCHINDUCTOR_CACHE_DIR=$ROOT/cache/torchinductor
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONPATH=$ROOT/Matrix-Game/Matrix-Game-2
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/matrix_native_controls.py" "$OUTPUT/probe.py"
+cp "$ROOT/Matrix-Game/Matrix-Game-2/inference.py" "$ROOT/Matrix-Game/Matrix-Game-2/configs/inference_yaml/inference_universal.yaml" "$OUTPUT/"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/matrix_native_controls.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/manifest.json

@@ -0,0 +1,50 @@
+{
+  "mode": "demo_actions",
+  "diagnostic_only": true,
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "upstream_revision": "71c3cd7f741311f8100f6cf9cde942b6c1378d11",
+  "weight_policy": "released raw distilled; original normalization; no optimizer or updates",
+  "expected_generator_parameters": 1619242304,
+  "expected_vae_parameters": 126892531,
+  "expected_clip_parameters": 1193014785,
+  "batch_size": 1,
+  "lora": null,
+  "lr": null,
+  "activation_checkpointing": false,
+  "steps": 3,
+  "shift": 5,
+  "cache_latents": 6,
+  "decoder": "official compiled max-autotune-no-cudagraphs",
+  "official_demo": "unchanged inference.py; supplied painting0000.png; seed0;150latents=597RGB",
+  "minecraft": {
+    "anchor": 699,
+    "native_latents": 9,
+    "native_rgb": 33,
+    "scored_slice": [
+      1,
+      25
+    ],
+    "targets": [
+      700,
+      724
+    ],
+    "recorded_control_alignment": "incoming700:724 plus9noops; small_mouse scales camera only",
+    "reference": "rawRGB with official center crop and frame_process"
+  },
+  "score_caveat": "Only recorded controls replay GT actions; other modes test stability, not benchmark accuracy",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "output_frames_per_case": 33,
+  "generator_requires_grad_parameters": 1619242304,
+  "inference_grad_policy": "unchanged native no_grad; parameter flags unmodified"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85114"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_demo_actions_job85112_task1/probe.py

@@ -0,0 +1,106 @@
+"""Official Matrix demo and native-image controls; all scores are diagnostic."""
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+import numpy as np
+import torch
+from PIL import Image
+from diagnose_matrix_game2 import save_result
+from protocol import minecraft_actions, read_frames, select_cases
+from worldmem_adapter import map_actions
+
+ROOT = Path(__file__).resolve().parents[1]
+REPO = ROOT / 'Matrix-Game/Matrix-Game-2'
+MODES = ['official_demo', 'demo_actions', 'neutral', 'recorded', 'small_mouse']
+
+def replay_controls(raw, mode, length):
+    keyboard, mouse = map_actions(minecraft_actions(raw), 1.0)
+    keys, turns = torch.zeros(length, 4), torch.zeros(length, 2)
+    if mode != 'neutral':
+        keys[:24], turns[:24] = keyboard, mouse * (0.1 if mode == 'small_mouse' else 1.0)
+    return {'keyboard_condition': keys, 'mouse_condition': turns}
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--mode', choices=MODES, required=True)
+    parser.add_argument('--data-root', type=Path, required=True)
+    parser.add_argument('--output', type=Path, required=True)
+    parser.add_argument('--dry-run', action='store_true')
+    args = parser.parse_args()
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    manifest = {'mode': args.mode, 'diagnostic_only': True, 'checkpoint': str(checkpoint),
+                'upstream_revision': '71c3cd7f741311f8100f6cf9cde942b6c1378d11',
+                'weight_policy': 'released raw distilled; original normalization; no optimizer or updates',
+                'expected_generator_parameters': 1619242304, 'expected_vae_parameters': 126892531,
+                'expected_clip_parameters': 1193014785, 'batch_size': 1, 'lora': None, 'lr': None,
+                'activation_checkpointing': False, 'steps': 3, 'shift': 5, 'cache_latents': 6,
+                'decoder': 'official compiled max-autotune-no-cudagraphs',
+                'official_demo': 'unchanged inference.py; supplied painting0000.png; seed0;150latents=597RGB',
+                'minecraft': {'anchor': 699, 'native_latents': 9, 'native_rgb': 33, 'scored_slice': [1, 25],
+                              'targets': [700, 724], 'recorded_control_alignment': 'incoming700:724 plus9noops; small_mouse scales camera only',
+                              'reference': 'rawRGB with official center crop and frame_process'},
+                'score_caveat': 'Only recorded controls replay GT actions; other modes test stability, not benchmark accuracy',
+                'cases': ([{'image': 'demo_images/universal/0000.png', 'seed': 0}] if args.mode == 'official_demo'
+                          else [{'index': r['global_test_index'], 'path': r['relative_path'],
+                                 'seed': 42000000 + int(r['global_test_index'])} for r in rows]),
+                'output_frames_per_case': 597 if args.mode == 'official_demo' else 33}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    os.chdir(REPO)
+    if args.mode == 'official_demo':
+        subprocess.run([sys.executable, 'inference.py', '--checkpoint_path', str(checkpoint / 'base_distilled_model/base_distill.safetensors'),
+                        '--pretrained_model_path', str(checkpoint), '--img_path', 'demo_images/universal/0000.png',
+                        '--config_path', 'configs/inference_yaml/inference_universal.yaml', '--seed', '0',
+                        '--num_output_frames', '150', '--output_folder', str(args.output)], check=True)
+        from torchvision.io import read_video
+        video, _, _ = read_video(str(args.output / 'demo.mp4'), pts_unit='sec')
+        assert len(video) == 597
+        Image.fromarray(torch.cat([video[t] for t in [0, 1, 2, 24, 100, 596]], dim=1).numpy()).save(args.output / 'samples.jpg')
+    else:
+        import inference as native
+        native.set_seed(0)
+        model_args = argparse.Namespace(config_path='configs/inference_yaml/inference_universal.yaml',
+            checkpoint_path=str(checkpoint / 'base_distilled_model/base_distill.safetensors'),
+            pretrained_model_path=str(checkpoint), num_output_frames=9, seed=0, img_path='', output_folder='')
+        engine = native.InteractiveGameInference(model_args)
+        manifest['generator_requires_grad_parameters'] = sum(p.numel() for p in engine.pipeline.generator.parameters() if p.requires_grad)
+        manifest['inference_grad_policy'] = 'unchanged native no_grad; parameter flags unmodified'
+        (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+        original_inference, native_controls = engine.pipeline.inference, native.Bench_actions_universal
+        captured = {}
+        def capture(**kwargs):
+            condition = kwargs['conditional_dict']
+            np.savez_compressed(case_dir / 'inputs.npz', initial_noise=kwargs['noise'].float().cpu().numpy(),
+                keyboard=condition['keyboard_cond'][0].float().cpu().numpy(), mouse=condition['mouse_cond'][0].float().cpu().numpy())
+            videos = original_inference(**kwargs)
+            captured['prediction'] = torch.cat(videos, dim=1)[0, 1:25].detach().float().cpu().add(1).div(2)
+            return videos
+        engine.pipeline.inference = capture
+        for row in rows:
+            case_dir = args.output / f"case{row['global_test_index']}"
+            case_dir.mkdir(exist_ok=True)
+            path = args.data_root / row['relative_path']
+            rgb = read_frames(path, torch.arange(699, 724))
+            images = [Image.fromarray(frame.mul(255).round().byte().permute(1, 2, 0).numpy()) for frame in rgb]
+            images[0].save(case_dir / 'input.png')
+            reference = torch.stack([engine.frame_process(engine._resizecrop(im, 352, 640)).add(1).div(2) for im in images[1:]])
+            with np.load(path.with_suffix('.npz')) as annotation:
+                controls = replay_controls(annotation['actions'][700:724], args.mode, 33)
+            native.Bench_actions_universal = native_controls if args.mode == 'demo_actions' else lambda length: controls
+            model_args.img_path, model_args.output_folder = str(case_dir / 'input.png'), str(case_dir)
+            engine.config['mode'] = 'universal'
+            native.set_seed(42000000 + int(row['global_test_index']))
+            engine.generate_videos()
+            save_result(case_dir, 'comparison', captured.pop('prediction'), reference)
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'complete': True}))
+    print('NATIVE_CONTROL_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/case0/comparison.json

@@ -0,0 +1,81 @@
+{
+  "case": "comparison",
+  "raw_rgb_psnr": [
+    29.00463104248047,
+    25.083839416503906,
+    5.1293206214904785,
+    7.163086891174316,
+    7.246057510375977,
+    7.313009262084961,
+    7.356295585632324,
+    7.3129777908325195,
+    7.402219772338867,
+    7.4208784103393555,
+    7.457307815551758,
+    7.469139099121094,
+    7.432450771331787,
+    7.4313764572143555,
+    7.441709518432617,
+    7.446588516235352,
+    7.415165901184082,
+    7.416438579559326,
+    7.437826156616211,
+    7.455380439758301,
+    7.600482940673828,
+    7.668927192687988,
+    7.7082390785217285,
+    7.726196765899658
+  ],
+  "raw_rgb_mse": [
+    0.0012575836153700948,
+    0.0031018152367323637,
+    0.3069502115249634,
+    0.19217252731323242,
+    0.18853598833084106,
+    0.1856517493724823,
+    0.1838105469942093,
+    0.18565309047698975,
+    0.18187710642814636,
+    0.1810973733663559,
+    0.17958466708660126,
+    0.1790960729122162,
+    0.18061546981334686,
+    0.18066012859344482,
+    0.18023081123828888,
+    0.18002846837043762,
+    0.18133574724197388,
+    0.1812826246023178,
+    0.1803920418024063,
+    0.17966437339782715,
+    0.1737607717514038,
+    0.1710437834262848,
+    0.16950249671936035,
+    0.16880306601524353
+  ],
+  "mean_rgb": [
+    0.71062171459198,
+    0.712151288986206,
+    0.20236869156360626,
+    0.3093110918998718,
+    0.31275486946105957,
+    0.316022664308548,
+    0.31764835119247437,
+    0.31517931818962097,
+    0.31939250230789185,
+    0.32033178210258484,
+    0.32227084040641785,
+    0.3230479955673218,
+    0.32040339708328247,
+    0.31977322697639465,
+    0.3199549615383148,
+    0.32007846236228943,
+    0.31954720616340637,
+    0.32010453939437866,
+    0.3214494287967682,
+    0.3225550353527069,
+    0.3296678364276886,
+    0.33288630843162537,
+    0.33478885889053345,
+    0.3357396125793457
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/case0/inputs.npz

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

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/case1/comparison.json

@@ -0,0 +1,81 @@
+{
+  "case": "comparison",
+  "raw_rgb_psnr": [
+    20.059900283813477,
+    18.597415924072266,
+    18.113113403320312,
+    17.338909149169922,
+    18.3916072845459,
+    18.688901901245117,
+    18.53306007385254,
+    18.357519149780273,
+    18.232343673706055,
+    18.186416625976562,
+    18.135757446289062,
+    18.113117218017578,
+    18.16193389892578,
+    18.166234970092773,
+    18.167783737182617,
+    18.165790557861328,
+    18.197973251342773,
+    18.205169677734375,
+    18.255691528320312,
+    18.27703285217285,
+    18.182619094848633,
+    18.162242889404297,
+    18.10346794128418,
+    18.056255340576172
+  ],
+  "raw_rgb_mse": [
+    0.009863022714853287,
+    0.013812059536576271,
+    0.015441467985510826,
+    0.018454786390066147,
+    0.01448235847055912,
+    0.013524146750569344,
+    0.014018258079886436,
+    0.014596477150917053,
+    0.015023309737443924,
+    0.015183024108409882,
+    0.015361170284450054,
+    0.01544145680963993,
+    0.015268862247467041,
+    0.015253743156790733,
+    0.015248307026922703,
+    0.015255306847393513,
+    0.015142681077122688,
+    0.015117603354156017,
+    0.01494276151061058,
+    0.014869509264826775,
+    0.015196304768323898,
+    0.015267775394022465,
+    0.0154758021235466,
+    0.01564495824277401
+  ],
+  "mean_rgb": [
+    0.7140963077545166,
+    0.7141072154045105,
+    0.7137575149536133,
+    0.7141113877296448,
+    0.7121320366859436,
+    0.7116583585739136,
+    0.7107756733894348,
+    0.7105647325515747,
+    0.7109081149101257,
+    0.711028516292572,
+    0.7102257609367371,
+    0.7101965546607971,
+    0.7101762294769287,
+    0.7103557586669922,
+    0.7101123332977295,
+    0.7103460431098938,
+    0.7076926827430725,
+    0.7064762115478516,
+    0.7051817178726196,
+    0.7046799063682556,
+    0.7081876993179321,
+    0.7094244956970215,
+    0.7100866436958313,
+    0.7104617953300476
+  ]
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/case1/inputs.npz

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

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/complete.json

@@ -0,0 +1 @@
+{"mode": "neutral", "complete": true}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/inference.py

@@ -0,0 +1,169 @@
+import os
+import argparse
+import torch
+import numpy as np
+
+from omegaconf import OmegaConf
+from torchvision.transforms import v2
+from diffusers.utils import load_image
+from einops import rearrange
+from pipeline import CausalInferencePipeline
+from wan.vae.wanx_vae import get_wanx_vae_wrapper
+from demo_utils.vae_block3 import VAEDecoderWrapper
+from utils.visualize import process_video
+from utils.misc import set_seed
+from utils.conditions import *
+from utils.wan_wrapper import WanDiffusionWrapper
+from safetensors.torch import load_file
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, default="configs/inference_yaml/inference_universal.yaml", help="Path to the config file")
+    parser.add_argument("--checkpoint_path", type=str, default="", help="Path to the checkpoint")
+    parser.add_argument("--img_path", type=str, default="demo_images/universal/0000.png", help="Path to the image")
+    parser.add_argument("--output_folder", type=str, default="outputs/", help="Output folder")
+    parser.add_argument("--num_output_frames", type=int, default=150,
+                        help="Number of output latent frames")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--pretrained_model_path", type=str, default="Matrix-Game-2.0", help="Path to the VAE model folder")
+    args = parser.parse_args()
+    return args
+
+class InteractiveGameInference:
+    def __init__(self, args):
+        self.args = args
+        self.device = torch.device("cuda")
+        self.weight_dtype = torch.bfloat16
+
+        self._init_config()
+        self._init_models()
+
+        self.frame_process = v2.Compose([
+            v2.Resize(size=(352, 640), antialias=True),
+            v2.ToTensor(),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+
+    def _init_config(self):
+        self.config = OmegaConf.load(self.args.config_path)
+
+    def _init_models(self):
+        # Initialize pipeline
+        generator = WanDiffusionWrapper(
+            **getattr(self.config, "model_kwargs", {}), is_causal=True)
+        current_vae_decoder = VAEDecoderWrapper()
+        vae_state_dict = torch.load(os.path.join(self.args.pretrained_model_path, "Wan2.1_VAE.pth"), map_location="cpu")
+        decoder_state_dict = {}
+        for key, value in vae_state_dict.items():
+            if 'decoder.' in key or 'conv2' in key:
+                decoder_state_dict[key] = value
+        current_vae_decoder.load_state_dict(decoder_state_dict)
+        current_vae_decoder.to(self.device, torch.float16)
+        current_vae_decoder.requires_grad_(False)
+        current_vae_decoder.eval()
+        current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")
+        pipeline = CausalInferencePipeline(self.config, generator=generator, vae_decoder=current_vae_decoder)
+        if self.args.checkpoint_path:
+            print("Loading Pretrained Model...")
+            state_dict = load_file(self.args.checkpoint_path)
+            pipeline.generator.load_state_dict(state_dict)
+
+        self.pipeline = pipeline.to(device=self.device, dtype=self.weight_dtype)
+        self.pipeline.vae_decoder.to(torch.float16)
+
+        vae = get_wanx_vae_wrapper(self.args.pretrained_model_path, torch.float16)
+        vae.requires_grad_(False)
+        vae.eval()
+        self.vae = vae.to(self.device, self.weight_dtype)
+
+    def _resizecrop(self, image, th, tw):
+        w, h = image.size
+        if h / w > th / tw:
+            new_w = int(w)
+            new_h = int(new_w * th / tw)
+        else:
+            new_h = int(h)
+            new_w = int(new_h * tw / th)
+        left = (w - new_w) / 2
+        top = (h - new_h) / 2
+        right = (w + new_w) / 2
+        bottom = (h + new_h) / 2
+        image = image.crop((left, top, right, bottom))
+        return image
+    
+    def generate_videos(self):
+        mode = self.config.pop('mode')
+        assert mode in ['universal', 'gta_drive', 'templerun']
+
+        image = load_image(self.args.img_path)
+        image = self._resizecrop(image, 352, 640)
+        image = self.frame_process(image)[None, :, None, :, :].to(dtype=self.weight_dtype, device=self.device)
+        # Encode the input image as the first latent
+        padding_video = torch.zeros_like(image).repeat(1, 1, 4 * (self.args.num_output_frames - 1), 1, 1)
+        img_cond = torch.concat([image, padding_video], dim=2)
+        tiler_kwargs={"tiled": True, "tile_size": [44, 80], "tile_stride": [23, 38]}
+        img_cond = self.vae.encode(img_cond, device=self.device, **tiler_kwargs).to(self.device)
+        mask_cond = torch.ones_like(img_cond)
+        mask_cond[:, :, 1:] = 0
+        cond_concat = torch.cat([mask_cond[:, :4], img_cond], dim=1) 
+        visual_context = self.vae.clip.encode_video(image)
+        sampled_noise = torch.randn(
+            [1, 16,self.args.num_output_frames, 44, 80], device=self.device, dtype=self.weight_dtype
+        )
+        num_frames = (self.args.num_output_frames - 1) * 4 + 1
+        
+        conditional_dict = {
+            "cond_concat": cond_concat.to(device=self.device, dtype=self.weight_dtype),
+            "visual_context": visual_context.to(device=self.device, dtype=self.weight_dtype)
+        }
+        
+        if mode == 'universal':
+            cond_data = Bench_actions_universal(num_frames)
+            mouse_condition = cond_data['mouse_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+            conditional_dict['mouse_cond'] = mouse_condition
+        elif mode == 'gta_drive':
+            cond_data = Bench_actions_gta_drive(num_frames)
+            mouse_condition = cond_data['mouse_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+            conditional_dict['mouse_cond'] = mouse_condition
+        else:
+            cond_data = Bench_actions_templerun(num_frames)
+        keyboard_condition = cond_data['keyboard_condition'].unsqueeze(0).to(device=self.device, dtype=self.weight_dtype)
+        conditional_dict['keyboard_cond'] = keyboard_condition
+        
+        with torch.no_grad():
+            videos = self.pipeline.inference(
+                noise=sampled_noise,
+                conditional_dict=conditional_dict,
+                return_latents=False,
+                mode=mode,
+                profile=False
+            )
+
+        videos_tensor = torch.cat(videos, dim=1)
+        videos = rearrange(videos_tensor, "B T C H W -> B T H W C")
+        videos = ((videos.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8)[0]
+        video = np.ascontiguousarray(videos)
+        mouse_icon = 'assets/images/mouse.png'
+        if mode != 'templerun':
+            config = (
+                keyboard_condition[0].float().cpu().numpy(),
+                mouse_condition[0].float().cpu().numpy()
+            )
+        else:
+            config = (
+                keyboard_condition[0].float().cpu().numpy()
+            )
+        process_video(video.astype(np.uint8), self.args.output_folder+f'/demo.mp4', config, mouse_icon, mouse_scale=0.1, process_icon=False, mode=mode)
+        process_video(video.astype(np.uint8), self.args.output_folder+f'/demo_icon.mp4', config, mouse_icon, mouse_scale=0.1, process_icon=True, mode=mode)
+        print("Done")
+
+def main():
+    """Main entry point for video generation."""
+    args = parse_args()
+    set_seed(args.seed)
+    os.makedirs(args.output_folder, exist_ok=True)
+    pipeline = InteractiveGameInference(args)
+    pipeline.generate_videos()
+
+if __name__ == "__main__":
+    main()

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/inference_universal.yaml

@@ -0,0 +1,21 @@
+denoising_step_list:
+- 1000
+- 666
+- 333
+warp_denoising_step: true
+ts_schedule: false
+mixed_precision: true
+seed: 42
+image_or_video_shape:
+- 1
+- 16
+- 15
+- 44
+- 80
+num_frame_per_block: 3
+context_noise: 0
+mode: universal
+causal: true
+model_kwargs:
+  timestep_shift: 5.0
+  model_config: configs/distilled_model/universal

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/launcher.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=matrix-game2-native-control
+#SBATCH --array=0-4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:1
+#SBATCH --time=02:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+PYTHON=/proj/cvl/users/x_fahkh2/envs/matrix-game2-eval/bin/python
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MODES=(official_demo demo_actions neutral recorded small_mouse)
+MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
+OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_native_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
+export TMPDIR=$ROOT/cache
+export TRITON_CACHE_DIR=$ROOT/cache
+export TORCHINDUCTOR_CACHE_DIR=$ROOT/cache/torchinductor
+export HF_HOME=$ROOT/cache/huggingface
+export TORCH_HOME=$ROOT/cache/torch
+export XDG_CACHE_HOME=$ROOT/cache
+export MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONPATH=$ROOT/Matrix-Game/Matrix-Game-2
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR" "$OUTPUT"
+"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$ROOT/deployment/matrix_native_controls.py" "$OUTPUT/probe.py"
+cp "$ROOT/Matrix-Game/Matrix-Game-2/inference.py" "$ROOT/Matrix-Game/Matrix-Game-2/configs/inference_yaml/inference_universal.yaml" "$OUTPUT/"
+srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/deployment/matrix_native_controls.py" \
+    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/manifest.json

@@ -0,0 +1,50 @@
+{
+  "mode": "neutral",
+  "diagnostic_only": true,
+  "checkpoint": "/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/checkpoints/matrix-game2",
+  "upstream_revision": "71c3cd7f741311f8100f6cf9cde942b6c1378d11",
+  "weight_policy": "released raw distilled; original normalization; no optimizer or updates",
+  "expected_generator_parameters": 1619242304,
+  "expected_vae_parameters": 126892531,
+  "expected_clip_parameters": 1193014785,
+  "batch_size": 1,
+  "lora": null,
+  "lr": null,
+  "activation_checkpointing": false,
+  "steps": 3,
+  "shift": 5,
+  "cache_latents": 6,
+  "decoder": "official compiled max-autotune-no-cudagraphs",
+  "official_demo": "unchanged inference.py; supplied painting0000.png; seed0;150latents=597RGB",
+  "minecraft": {
+    "anchor": 699,
+    "native_latents": 9,
+    "native_rgb": 33,
+    "scored_slice": [
+      1,
+      25
+    ],
+    "targets": [
+      700,
+      724
+    ],
+    "recorded_control_alignment": "incoming700:724 plus9noops; small_mouse scales camera only",
+    "reference": "rawRGB with official center crop and frame_process"
+  },
+  "score_caveat": "Only recorded controls replay GT actions; other modes test stability, not benchmark accuracy",
+  "cases": [
+    {
+      "index": "0",
+      "path": "test/val_plains/2/000034.mp4",
+      "seed": 42000000
+    },
+    {
+      "index": "1",
+      "path": "test/val_desert/1/000029.mp4",
+      "seed": 42000001
+    }
+  ],
+  "output_frames_per_case": 33,
+  "generator_requires_grad_parameters": 1619242304,
+  "inference_grad_policy": "unchanged native no_grad; parameter flags unmodified"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/node.json

@@ -0,0 +1,7 @@
+{
+  "devices": [
+    "NVIDIA H200"
+  ],
+  "torch": "2.7.0+cu126",
+  "slurm_job_id": "85115"
+}

results/worldmem-baseline-evals/matrix-game2/diagnostic_native_neutral_job85112_task2/probe.py

@@ -0,0 +1,106 @@
+"""Official Matrix demo and native-image controls; all scores are diagnostic."""
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+import numpy as np
+import torch
+from PIL import Image
+from diagnose_matrix_game2 import save_result
+from protocol import minecraft_actions, read_frames, select_cases
+from worldmem_adapter import map_actions
+
+ROOT = Path(__file__).resolve().parents[1]
+REPO = ROOT / 'Matrix-Game/Matrix-Game-2'
+MODES = ['official_demo', 'demo_actions', 'neutral', 'recorded', 'small_mouse']
+
+def replay_controls(raw, mode, length):
+    keyboard, mouse = map_actions(minecraft_actions(raw), 1.0)
+    keys, turns = torch.zeros(length, 4), torch.zeros(length, 2)
+    if mode != 'neutral':
+        keys[:24], turns[:24] = keyboard, mouse * (0.1 if mode == 'small_mouse' else 1.0)
+    return {'keyboard_condition': keys, 'mouse_condition': turns}
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--mode', choices=MODES, required=True)
+    parser.add_argument('--data-root', type=Path, required=True)
+    parser.add_argument('--output', type=Path, required=True)
+    parser.add_argument('--dry-run', action='store_true')
+    args = parser.parse_args()
+    checkpoint = ROOT / 'checkpoints/matrix-game2'
+    rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2)
+    manifest = {'mode': args.mode, 'diagnostic_only': True, 'checkpoint': str(checkpoint),
+                'upstream_revision': '71c3cd7f741311f8100f6cf9cde942b6c1378d11',
+                'weight_policy': 'released raw distilled; original normalization; no optimizer or updates',
+                'expected_generator_parameters': 1619242304, 'expected_vae_parameters': 126892531,
+                'expected_clip_parameters': 1193014785, 'batch_size': 1, 'lora': None, 'lr': None,
+                'activation_checkpointing': False, 'steps': 3, 'shift': 5, 'cache_latents': 6,
+                'decoder': 'official compiled max-autotune-no-cudagraphs',
+                'official_demo': 'unchanged inference.py; supplied painting0000.png; seed0;150latents=597RGB',
+                'minecraft': {'anchor': 699, 'native_latents': 9, 'native_rgb': 33, 'scored_slice': [1, 25],
+                              'targets': [700, 724], 'recorded_control_alignment': 'incoming700:724 plus9noops; small_mouse scales camera only',
+                              'reference': 'rawRGB with official center crop and frame_process'},
+                'score_caveat': 'Only recorded controls replay GT actions; other modes test stability, not benchmark accuracy',
+                'cases': ([{'image': 'demo_images/universal/0000.png', 'seed': 0}] if args.mode == 'official_demo'
+                          else [{'index': r['global_test_index'], 'path': r['relative_path'],
+                                 'seed': 42000000 + int(r['global_test_index'])} for r in rows]),
+                'output_frames_per_case': 597 if args.mode == 'official_demo' else 33}
+    if args.dry_run:
+        print(json.dumps(manifest, indent=2))
+        return
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+    os.chdir(REPO)
+    if args.mode == 'official_demo':
+        subprocess.run([sys.executable, 'inference.py', '--checkpoint_path', str(checkpoint / 'base_distilled_model/base_distill.safetensors'),
+                        '--pretrained_model_path', str(checkpoint), '--img_path', 'demo_images/universal/0000.png',
+                        '--config_path', 'configs/inference_yaml/inference_universal.yaml', '--seed', '0',
+                        '--num_output_frames', '150', '--output_folder', str(args.output)], check=True)
+        from torchvision.io import read_video
+        video, _, _ = read_video(str(args.output / 'demo.mp4'), pts_unit='sec')
+        assert len(video) == 597
+        Image.fromarray(torch.cat([video[t] for t in [0, 1, 2, 24, 100, 596]], dim=1).numpy()).save(args.output / 'samples.jpg')
+    else:
+        import inference as native
+        native.set_seed(0)
+        model_args = argparse.Namespace(config_path='configs/inference_yaml/inference_universal.yaml',
+            checkpoint_path=str(checkpoint / 'base_distilled_model/base_distill.safetensors'),
+            pretrained_model_path=str(checkpoint), num_output_frames=9, seed=0, img_path='', output_folder='')
+        engine = native.InteractiveGameInference(model_args)
+        manifest['generator_requires_grad_parameters'] = sum(p.numel() for p in engine.pipeline.generator.parameters() if p.requires_grad)
+        manifest['inference_grad_policy'] = 'unchanged native no_grad; parameter flags unmodified'
+        (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
+        original_inference, native_controls = engine.pipeline.inference, native.Bench_actions_universal
+        captured = {}
+        def capture(**kwargs):
+            condition = kwargs['conditional_dict']
+            np.savez_compressed(case_dir / 'inputs.npz', initial_noise=kwargs['noise'].float().cpu().numpy(),
+                keyboard=condition['keyboard_cond'][0].float().cpu().numpy(), mouse=condition['mouse_cond'][0].float().cpu().numpy())
+            videos = original_inference(**kwargs)
+            captured['prediction'] = torch.cat(videos, dim=1)[0, 1:25].detach().float().cpu().add(1).div(2)
+            return videos
+        engine.pipeline.inference = capture
+        for row in rows:
+            case_dir = args.output / f"case{row['global_test_index']}"
+            case_dir.mkdir(exist_ok=True)
+            path = args.data_root / row['relative_path']
+            rgb = read_frames(path, torch.arange(699, 724))
+            images = [Image.fromarray(frame.mul(255).round().byte().permute(1, 2, 0).numpy()) for frame in rgb]
+            images[0].save(case_dir / 'input.png')
+            reference = torch.stack([engine.frame_process(engine._resizecrop(im, 352, 640)).add(1).div(2) for im in images[1:]])
+            with np.load(path.with_suffix('.npz')) as annotation:
+                controls = replay_controls(annotation['actions'][700:724], args.mode, 33)
+            native.Bench_actions_universal = native_controls if args.mode == 'demo_actions' else lambda length: controls
+            model_args.img_path, model_args.output_folder = str(case_dir / 'input.png'), str(case_dir)
+            engine.config['mode'] = 'universal'
+            native.set_seed(42000000 + int(row['global_test_index']))
+            engine.generate_videos()
+            save_result(case_dir, 'comparison', captured.pop('prediction'), reference)
+    (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'complete': True}))
+    print('NATIVE_CONTROL_COMPLETE', flush=True)
+
+if __name__ == '__main__':
+    main()