Spatial-MMCoT

Research code for studying whether action-conditioned multimodal generation can act as an explicit spatial world model and transfer to spatial reasoning.

Third-batch GRPO experiments use OURS/config/rl/RLv3/, with shared rollout and validation queues, per-query evaluation recovery, and SIQ one-hour A100×4 attempts. See training objectives.

Layout

• AGENTS.md: research direction and repository-wide contracts.
• OURS/: installable final-library code for the independent UniAR backbone and reusable shared/SFT/OPD/RL namespaces.
• OURS/src/uniar/: checkpoint-compatible Qwen3-VL/BSQ backbone, visual head, token protocol, AR rollout, and RGB decoder.
• OURS/src/shared/: the only cross-training-module layer for distributions, EMA, model construction, supervised replay, answer scoring, and artifact primitives.
• OURS/src/sft/: canonical cache reader, interleaved text/bit/Soft-BSQ/Information-Gain objectives, symmetric-JSD EMA Necessity distillation, distributed trainer, logging, and atomic resume.
• OURS/src/opd/: single-/multi-GPU full-trajectory on-policy distillation, optional mass-guarded text top-K JSD, GT replay with branch-gradient diagnostics, and persistent shared-queue rollout/validation workers. Validation saves complete trajectories and TensorBoard metrics.
• OURS/config/sft/: complete, strict YAML configurations for SFT runs.
• OURS/config/opd/: strict OPD v0 routing/replay configurations and versioned dataset experiments.
• OURS/src/rl/: config-driven one- or multi-GPU multimodal GRPO v0 with stochastic full-trajectory rollout, strict Visual-Jigsaw A–D verification, optional detached four-view IG reward, factor-wise visual PPO, and atomic resume.
• OURS/config/rl/: strict GRPO v0 baseline and smoke configurations; image reward is fixed to literal zero with no backend.
• OURS-Bagel/: canonical BAGEL-family framework with strict YAML SFT and LoRA OPD for text categorical targets and recorded reverse-path flow targets, necessity/sufficiency, EMA, frozen-reference anchoring and GT replay, with fixed-batch dynamic rollout/backward queues and per-query recovery. Shared native adapters live in bagel_mmcot/shared; rl remains reserved.
• OURS-LatentUM/: canonical LatentUM-family framework. Its migration stage owns strict runtime/SFT YAML, typed protocol rendering, lazy upstream model access, paper-joint text/visual SFT, native interleaved VLMEvalKit adaptation, tests, and reserved opd/rl boundaries.
• experiments/src/: dataset preparation, inference, evaluation, controllers, and historical SFT workflows.
• experiments/src/zebra_baselines/: native-protocol adapters for released Bagel-Zebra-CoT and Anole-Zebra-CoT checkpoints.
• experiments/external/: pinned external research-code checkouts and papers.
• projects/UniAR/: read-only upstream UniAR source reference.
• projects/ThinkMorph-edited/: pinned BAGEL architecture/data/FSDP execution dependency used by OURS-Bagel; its paper-aligned contract is in docs/thinkmorph.md.
• projects/LatentUM/: pinned model/inference dependency used by OURS-LatentUM; its migration audit and implementation gates are in docs/latentum.md.
• ThinkMorph-OPD/: legacy native ThinkMorph-7B OPD prototype, pending migration into OURS-Bagel; it contains the heterogeneous carrier tracer, and V→T-only LoRA/EMA OPD with C11/C10 masking, DeepSeek failure routing, full native text+flow GT replay, distributed updates, and resume.
• projects/VLMEvalKit_Thinkmorph_llm/: canonical external evaluation fork with the UniAR interleaved adapter, upstream verifier, and DeepSeek V4 Flash fallback judge.
• models/UniAR-SFT/: canonical released UniAR SFT checkpoint.
• models/Bagel-7B-MoT/: pinned EMA initialization for the ThinkMorph baseline.
• models/Qwen3-VL-8B-Instruct/: pinned Qwen3-VL comparison checkpoint.
• models/Qwen3.8-27B/: pinned Qwen3.8-27B multimodal baseline checkpoint.
• data/Zebra-CoT: shared-storage link to the full Zebra-CoT download; its seven Graph shards are locally repaired with official reasoning_image_10..19 columns.
• data/Zebra-CoT-samples: inspectable 16-trajectory raw/prepared overfit fixture.
• data/Zebra-CoT-small: materialized 1% raw-Parquet development subset (1,825 rows across all 18 configs).
• data/Zebra-CoT-{Visual-Jigsaw,Chess,Maze}: physically materialized full single-task Parquet subsets with answer audits, group-safe 90/5/5 splits, and zero-copy UniAR-960 cache views for controlled SFT.
• data/Zebra-CoT-Visual-Jigsaw/UniAR-960-cache-deltaai: local rebuilt packed cache used on DeltaAI; the migrated zero-copy view retains unavailable old /mnt/ultracube paths and is provenance-only on this cluster.
• data/Zebra-CoT-UniAR-960: shared packed-BSQ cache for the full dynamic-resolution SFT corpus.
• experiments/sites/: independent self-contained Sites projects, copied assets, and publication archives; deployed sites never read assets from experiments/output/.
• experiments/output/: ignored model-first artifacts under <owner>/<dataset-or-topic>/<version>/{runs,sweeps,smoke,archive}/<task>/<timestamp_name>; shared and ops own non-model tasks. See artifact layout and Bucket safeguards.
• data/ThinkMorph/ThinkMorph: train-only canonical UniAR cache built from all 24,990 public ThinkMorph SFT rows; see docs/thinkmorph.md.
• data/ThinkMorph/ThinkMorph-IDVal5: fixed image-grouped 95/5 development view (23,741 train / 1,249 validation), sharing the original BSQ shards. Rank-64 2/4/8-epoch selection uses this view; final full-data training returns to the original all-train cache. See data/ThinkMorph/README.md.
• data/ThinkMorph/Bagel-SFT: official-format PNG/Parquet corpus for the paper-aligned BAGEL/ThinkMorph baseline.
• data/VLMEvalKit/BLINK_Jigsaw.tsv: checksum-verified 150-row external BLINK-Jigsaw evaluation set.

Environment

The shared cluster workspace is /work/nvme/bgjw/tcheng2: this repository is MMCoT/, applications live in apps/, and default caches/temp files in .cache/. Both clusters load the private sibling .env from their shell startup files; credentials are exported environment variables, never public project files. The local artifact tree may be sparse; retained large checkpoints are archived in hf://buckets/EmergentTwilight/MMCoT. Hydrate a complete checkpoint before resuming it; see archive safeguards.

MAMBA=/work/nvme/bgjw/tcheng2/apps/miniforge3-deltaai/bin/mamba
$MAMBA run -n UniAR python -m pip install -e 'OURS[train,decoder]'

On DeltaAI, use this Miniforge directly without loading cluster modules. The GH200 stack is PyTorch 2.7.1+cu128, torchvision 0.22.1, qwen-vl-utils 0.0.14, and FlashAttention 2.8.3.post1 built for SM90 against the NVIDIA HPC SDK CUDA 12.9. See docs/uniar.md for the complete setup. Submit GPU work through Slurm and let it assign devices; do not set CUDA_VISIBLE_DEVICES manually. Create new Slurm runs with the submission-time snapshot wrapper documented in docs/submission_snapshots.md; queued jobs import frozen _submission/source rather than the mutable working tree. UniAR SFT v1 configs write below experiments/output/uniar-sft/<dataset>/v1/; unversioned experiments use the explicit unversioned directory. BAGEL and LatentUM use their actual base-model slug with the same layout. Historical locations are indexed by experiments/output/layout-map.json. Every *-interactive run must be prepared with cluster_jobs.submit --defer --resumable and handed to the system-level siq app; direct interactive submission is rejected. SFT and autonomous evaluation reuse one output directory across all TIMEOUT retries. The integration and recovery contracts are documented in the same file. Chain interactive work with SIQ Task-ID dependencies so downstream evaluation remains outside the Dispatch FIFO until training completes successfully. Distributed SFT checks actual replica equality at startup/backward/checkpoint boundaries. Its NamedTuple objective output is part of the DDP contract; see docs/zebra_cot.md and the four-process regression in OURS/tests/test_sft_distributed_control.py. Use PYTHONPATH="$PWD/OURS/src" for the standalone library. Historical workflows use PYTHONPATH="$PWD/OURS/src:$PWD/experiments/src", with the owned backbone first. The upstream checkout is only a source reference; model weights live under models/. ThinkMorph/BAGEL work uses the separate ThinkMorph environment and canonical OURS-Bagel interface; projects/ThinkMorph-edited is its pinned execution dependency. Do not install that Transformers 4.49 runtime into UniAR. LatentUM work uses private LatentUM environments on both Delta and DeltaAI, with Torch 2.8.0+cu129 and FlashAttention 2.8.3. The CUDA 12.9 index provides both architecture wheels. Read docs/latentum.md for the environment contract, GPU verification status, paper-joint SFT, and native image/chat protocol. LatentUM SFT supports full tuning or two-branch PEFT LoRA; its ThinkMorph base recipes use the author's 0.15 training-only visual-context code noise. Optional absolute/relative IG is additive to that native SFT objective. latentum_sft.ig_sweep runs a resumable calibration/search bundle with isolated fixed-batch component-gradient norms, cosines and signed projection diagnostics; see docs/latentum.md before selecting or interpreting a search point. LatentUM constant-LR DDP downsizing can retain a committed model/Adam/data cursor through latentum_sft.prepare_continuation; prepare the parent-linked child before SIQ enqueue and follow the RNG-mapping constraints in docs/latentum.md. Epoch-end validation is enabled in the canonical configs; sft validate-epochs backfills completed-run validation into the original TensorBoard in step order with explicit checkpoint-step provenance. See docs/latentum.md. LatentUM benchmark evaluation uses the same private environment and either the native Base protocol or saved SFT protocol, with greedy text/latent rollouts, VSP official scoring, extraction-only LLM fallback and dynamic per-query resume; see docs/latentum.md for submission. VLMEvalKit UniAR evaluation uses ThinkMorph-Eval with Torch 2.7.1, Transformers 4.57, PEFT 0.20, and FlashAttention 2.8.3; both DeltaAI and Delta have architecture-native environments. See docs/thinkmorph.md. Historical Qwen3.5/Qwen3.8 CV-Bench inference uses the shared RL4SI-Qwen35-SFT environment described in docs/cv_bench.md. DeltaAI Qwen3.8 vLLM baselines use a separate Qwen38-vLLM environment because vLLM 0.28.0 requires Torch 2.13/Transformers 5.x and must not modify UniAR:

$MAMBA create -n Qwen38-vLLM python=3.12 pip -y
$MAMBA run -n Qwen38-vLLM python -m pip install \
  vllm==0.28.0 pyarrow pillow

Inspect the configured small-data provenance on CPU, then launch score probes or training on an assigned GPU:

PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m sft inspect \
  --config OURS/config/sft/zebra_ig_lora.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m sft probe \
  --config OURS/config/sft/zebra_ig_lora.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR torchrun --nproc-per-node=2 \
  -m sft train --config OURS/config/sft/zebra_ig_lora.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd inspect \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd train \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd audit-rgb \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml --samples 1
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m rl inspect \
  --config OURS/config/rl/zebra_visual_jigsaw_grpo_fmt_acc.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m rl train \
  --config OURS/config/rl/zebra_visual_jigsaw_grpo_fmt_acc.yaml
# Same entry point; trainer.world_size: 4 auto-spawns four local workers.
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m rl train \
  --config OURS/config/rl/zebra_visual_jigsaw_grpo_fmt_acc_exploratory_128_4gpu.yaml

Each GRPO v0 optimizer update uses only its freshly collected prompt groups and regularizes against a frozen run-start reference with full-distribution symmetric JSD. Generated grid sides are parsed only after the student emits <|vision_start|> and must lie in 1–30; RGB-decoder/scorer-based visual diagnostics remain offline/shadow analyses of saved artifacts. trainer.world_size is part of the strict config; multi-GPU workers shard one group's trajectories, reconstruct its global reward and advantage, then sum globally normalized gradients before one shared step. evaluation.rollouts_per_sample can reduce validation sampling independently of the training group size; those trajectories are scored without advantages.

The controlled Zebra-CoT-small one-epoch pair is configured by zebra_small_standard_1epoch.yaml and zebra_small_information_gain_1epoch.yaml under OURS/config/sft/. The current five-run Visual-Jigsaw BitFlip control matrix uses zebra_visual_jigsaw_control_*_bitflip_2gpu_1epoch.yaml. Counterfactual IG/EMA views are virtual-batch chunked; single-GPU execution may re-chunk a frozen logical update with trainer.execution_* fields without changing its rows or BitFlip realization. Validation is shard-local and batched. Rollout BSQ artifacts are bit-packed with backward-compatible readers.

Canonical documentation

• docs/uniar.md: UniAR setup, checkpoint, and runtime constraints.
• docs/cv_bench.md: native understanding baselines and scoring.
• docs/zebra_cot.md: Zebra-CoT audit, typed trajectory, joint SFT, held-out rollout, and causal visual-context ablation workflow.
• docs/modeling.md: model-independent AR/FM SFT-IG, OPD (discrete JSD/KL, flow velocity matching, EMA teacher and frozen reference), and RL method; implementation details and status are isolated in the final chapter.
• docs/research/20260824-122930_official_zebra_models_membership_and_cvbench_probe.md: released-checkpoint membership audit and sampled external-model results.
• docs/novel_view_sft.md: posed multi-view ingestion, camera semantics, QA generation, curriculum, and novel-view SFT evaluation.
• docs/visual_quality.md: offline and shadow visual-quality evaluation for saved rollout artifacts.
• docs/artifacts.md: output paths and preservation policy.
• docs/submission_snapshots.md: immutable queued-job source, launcher, and continuation contract.