SGE-Goal — LoRA checkpoints

LoRA adapters for SGE-Goal, from Goal Images Are Control Targets: Spatially-Grounded Goal-Image Synthesis for Robotic Manipulation (WACV 2027, arXiv:2610.24885).

SGE-Goal synthesizes a goal image for a goal-conditioned manipulation policy in three steps: a training-free initial edit-region mask M_i (Step 1), a fine-tuned final edit-region mask M_f (Step 2) and the fine-tuned goal image Î_e (Step 3). This repository holds the Step 2 and Step 3 adapters for both base editors and all three datasets, plus the mask-ablation adapters.

Code, data pipeline, evaluator and usage instructions: https://github.com/Riddhi-Chatterjee/SGE-Goal

Contents

Path	Base editor	Dataset	Stage	Epoch
flux1-kontext-dev/bridge/step2_final_mask/	FLUX.1 Kontext [dev]	BridgeDataV2	Step 2 · final mask	38
flux1-kontext-dev/calvin/step2_final_mask/	FLUX.1 Kontext [dev]	CALVIN-ABCD	Step 2 · final mask	22
flux1-kontext-dev/libero/step2_final_mask/	FLUX.1 Kontext [dev]	LIBERO-10/90	Step 2 · final mask	56
flux1-kontext-dev/bridge/step3_goal_image/	FLUX.1 Kontext [dev]	BridgeDataV2	Step 3 · goal image	18
flux1-kontext-dev/calvin/step3_goal_image/	FLUX.1 Kontext [dev]	CALVIN-ABCD	Step 3 · goal image	18
flux1-kontext-dev/libero/step3_goal_image/	FLUX.1 Kontext [dev]	LIBERO-10/90	Step 3 · goal image	64
flux2-klein/bridge/step2_final_mask/	FLUX.2 klein	BridgeDataV2	Step 2 · final mask	17
flux2-klein/calvin/step2_final_mask/	FLUX.2 klein	CALVIN-ABCD	Step 2 · final mask	22
flux2-klein/libero/step2_final_mask/	FLUX.2 klein	LIBERO-10/90	Step 2 · final mask	59
flux2-klein/bridge/step3_goal_image/	FLUX.2 klein	BridgeDataV2	Step 3 · goal image	9
flux2-klein/calvin/step3_goal_image/	FLUX.2 klein	CALVIN-ABCD	Step 3 · goal image	15
flux2-klein/libero/step3_goal_image/	FLUX.2 klein	LIBERO-10/90	Step 3 · goal image	71
ablations/flux1-kontext-dev/bridge/wo_final_mask/	FLUX.1 Kontext [dev]	BridgeDataV2	Step 3, without M_f	9
ablations/flux1-kontext-dev/bridge/wo_masks/	FLUX.1 Kontext [dev]	BridgeDataV2	Step 3, without M_i, M_f	13
ablations/flux2-klein/bridge/wo_final_mask/	FLUX.2 klein	BridgeDataV2	Step 3, without M_f	9
ablations/flux2-klein/bridge/wo_masks/	FLUX.2 klein	BridgeDataV2	Step 3, without M_i, M_f	10

Octo-ft: Octo fine-tuned on CALVIN (octo-ft/calvin/)

The third CALVIN policy of the paper's supplement (Sec. A): Octo-Base fine-tuned on real CALVIN task_D_D training frames (1,214 language-annotated windows, ~64k frames), used frozen and conditioned on the goal image only to evaluate goal-image generators. It never saw any generator's output.

Path	Model	Data	Steps
octo-ft/calvin/	Octo-Base, all parameters fine-tuned	CALVIN task_D_D (training split)	30,000 (final checkpoint)

Contents: the Orbax checkpoint (30000/), config.json, dataset_statistics.json (action un-normalisation), example_batch.msgpack and finetune_config.json. Recipe: batch 32, AdamW (wd 0.01), cosine LR peak 3e-4 with 1k warm-up, 2-step observation window, 4-step action chunk, CALVIN relative actions, goal image or instruction per training sample (p = 0.5), goal image augmented independently of the observation. Training and evaluation code: downstream/calvin/octo_ft/ in the GitHub repository.

from octo.model.octo_model import OctoModel   # Octo at commit 241fb35 + our patch (see the repo)
model = OctoModel.load_pretrained("/path/to/octo-ft/calvin")   # loads the step-30000 checkpoint

The model emits CALVIN rel_actions directly (no rescaling). Evaluate with the goal image only (OCTO_TASK_MODE=goal_only in the repo's wrapper): given the instruction as well, it largely ignores the goal image.

Each SGE-Goal directory holds the adapter weights (pytorch_lora_weights.safetensors, or transformer_lora.safetensors for the FLUX.2 klein Step 2 adapters) and, where available, the trainer's state file recording the epoch and step. Optimizer and scheduler states are not included. SHA256SUMS lists the checksum of every weight file.

Training: LoRA rank 16, α 16, 8-bit AdamW, learning rate 1e-4 (constant, 200 warm-up steps), batch size 160, bf16; 256 px for BridgeDataV2 and CALVIN, 128 px for LIBERO. Step 2 adds a channel-consistency loss (λ = 25) and Step 3 a static-scene-consistency loss (λ = 7), both applied for σ < 0.35.

Usage

huggingface-cli download RiddhiCh/SGE-Goal --local-dir sge-goal-checkpoints
sha256sum -c sge-goal-checkpoints/SHA256SUMS --ignore-missing   # run from inside the folder

Then point the pipeline drivers in the code repository at the matching directories, e.g. for FLUX.2 klein on BridgeDataV2:

cd sge_goal/pipelines/flux2_klein
python run_pipeline.py --data-dir /abs/path/to/test_data \
  --fer-checkpoint-dir /abs/path/sge-goal-checkpoints/flux2-klein/bridge/step2_final_mask \
  --ie-checkpoint-dir  /abs/path/sge-goal-checkpoints/flux2-klein/bridge/step3_goal_image

Licence

The adapters are released by the authors for research use. They are fine-tuned from FLUX.1 Kontext [dev] and FLUX.2 klein; using them requires the base model and therefore complying with each base model's licence (FLUX.1 [dev] is under the FLUX.1 [dev] Non-Commercial License). The licences of the training datasets (BridgeData V2, CALVIN, LIBERO) also apply.

Citation

If you use these checkpoints, please cite our paper:

@article{chatterjee2026goalimages,
  title   = {Goal Images Are Control Targets: Spatially-Grounded Goal-Image Synthesis for Robotic Manipulation},
  author  = {Chatterjee, Riddhi and Shrirao, Ritish and Mehrotra, Kushagra and Gopalakrishnan, Viswanath},
  journal = {arXiv preprint arXiv:2610.24885},
  year    = {2026}
}