GeoUG-ColdStart

Project page · Paper

The supervised cold-start checkpoint for GeoUG spatial-policy RL.

GeoUG-ColdStart is the supervised initialization checkpoint for spatial-policy reinforcement learning. It provides the starting point for training GeoUG-Spatial.

Property	Value
Stage	Interleaved reasoning supervised fine-tuning
Weights	EMA cold-start checkpoint
Precision	Original BF16 weights
Downstream model	GeoUG-Spatial after spatial-policy RL
Bundled transform	Matches reasoning SFT and the spatial RL recipe

Download

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="rayruiyang/GeoUG-ColdStart",
    local_dir="ckpt/geoug-coldstart",
)

These are custom BAGEL/MoT weights. Use the GeoUG runtime rather than a generic AutoModel.from_pretrained call. The weight repository contains the checkpoint assets; inference and training code live in the GeoUG code checkout.

Verify the download:

(cd ckpt/geoug-coldstart && sha256sum -c checksums.sha256)

Initialize RL training

Use the GeoUG checkout's separate RL environment, described in docs/RL_TRAINING.md. After downloading this model and train_geoug_rl.parquet, restore the public training layout:

python scripts/pack_data_parquet.py unpack-rl \
  --input /absolute/path/to/train_geoug_rl.parquet \
  --output /absolute/path/to/rl_train

export MODEL_PATH="$PWD/ckpt/geoug-coldstart"
export RL_TRANSFORM_CONFIG="$MODEL_PATH/transform_config.yaml"
export RL_DATASET=/absolute/path/to/rl_train
export DEPTH_CACHE_DIR=/absolute/path/to/rl_train/depth-cache

The parquet includes training images, camera metadata, and generated-depth caches. Prepare your own validation/evaluation inputs as required by the RL recipe; the internal RL evaluation subset is not distributed.

The transform uses multiscale disparity normalization with c: [2.0, 8.0, 32.0]. The RL loader uses ema.safetensors, matching the initialization used by the released recipe.

Scope and interpretation

This repository provides the initialization checkpoint, not the final adaptive policy. Use GeoUG-Spatial for the final model. Its reported RL benchmark scores must not be attributed to GeoUG-ColdStart.

The cold start model is released here. The cold start training data is a separate, pending release; it is not included in this repository or in the 800-sample RL training parquet.

Model family

Checkpoint	Intended use
GeoUG-Spatial	Final adaptive spatial reasoning policy
GeoUG-Depth	Standalone metric-depth generation
GeoUG-ColdStart	SFT initialization for spatial-policy RL

Package and integrity

File	Contents
ema.safetensors	BF16 EMA model weights, approximately 27.21 GiB
ae.safetensors	VAE weights
llm_config.json, vit_config.json, config.json, generation_config.json	Model and generation configuration
tokenizer.json, tokenizer_config.json, vocab.json, merges.txt	Tokenizer assets
transform_config.yaml	Image/depth preprocessing and normalization
checksums.sha256	SHA-256 checksums for the package files

Download the complete repository. The release uses EMA weights and does not require a separate non-EMA checkpoint or a shard index.

Checkpoint provenance

This is the EMA checkpoint at step 1500 of the v6 interleaved reasoning SFT run. The original SFT stage starts from the v5 step-1000 initialization. These training identifiers are retained for experiment traceability.

License and attribution

GeoUG builds on the BAGEL/MoT architecture and its pretrained components. This checkpoint accompanies the mixed-license GeoUG code release. Consult LICENSE and THIRD_PARTY_NOTICES.md in the GeoUG code checkout for the applicable component terms, including the CC BY-NC restrictions on the listed BAGEL/DiT-derived code files. This card does not assign a new blanket license to upstream components. The data repositories document their own source and usage terms separately.

Contact

For questions about this release, contact rayyang0116@connect.hku.hk.