# Spatial Ssrl Spatial Understanding Rl > Improve spatial reasoning in vision-language models through five automatically-formulated pretext tasks (patch shuffling, flipping, inpainting, depth ordering, 3D position) that require zero human annotation, using only RGB/RGB-D images and RL optimization for 4% accuracy gains. - Skill: `adu2021/spatial-ssrl-spatial-understanding-rl` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/spatial-ssrl-spatial-understanding-rl` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/spatial-ssrl-spatial-understanding-rl/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/spatial-ssrl-spatial-understanding-rl --- # Title: Train Spatial Reasoning Through Vision-Only Pretext Tasks and RL Vision-language models struggle with spatial relationships despite being trained on billions of image-text pairs. Spatial-SSRL fixes this by creating five pretext tasks that derive supervisory signals directly from image structure—no human annotation needed. Tasks range from simple (reorder shuffled patches) to geometric (predict 3D positions). The model learns through Group Relative Policy Optimization on automatically-generated question-answer pairs from existing datasets. The key insight is that images contain all necessary geometry; you just need the right pretext tasks to make it explicit. ## Core Concept **Self-Supervised Spatial Learning**: - **Five Pretext Tasks**: Shuffled patch reordering, flipped patch recognition, patch inpainting, regional depth ordering, relative position prediction - **Zero-Human-Annotation**: All ground truth derived from image structure or depth maps - **GRPO Optimization**: Group Relative Policy Optimization on structured reasoning - **Dual-Modality Support**: RGB-only tasks and RGB-D (depth-dependent) tasks - **Automatic Dataset Generation**: 81K QA pairs from existing datasets (COCO, DIODE, MegaDepth) ## Architecture Overview - **Pretext Task Set**: Two categories (RGB-free, depth-based) with automatic question generation - **Training Data**: Spatial-SSRL-81k dataset with no manual annotation - **GRPO Framework**: Reward combination (accuracy 90% + format 10%) using reasoning-style outputs - **Cold-Start Phase**: Brief supervised fine-tuning before RL optimization - **Target Models**: Applied to Qwen2.5-VL (3B, 7B variants) ## Implementation Steps **1. Define Five Pretext Tasks with Automatic Question Generation** Create task-specific question-answer generation from images. ```python class SpatialPretext: @staticmethod def shuffled_patch_reordering(image, patch_size=32, num_patches=16): # Divide image into patches patches = image.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size) # Randomly shuffle perm = torch.randperm(num_patches) shuffled = patches[perm] # Question: restore original order question = "The image patches are shuffled. Reorder them to restore the original." answer = f"The correct order is: {list(range(num_patches))}" return question, answer @staticmethod def flipped_patch_recognition(image, patch_size=32): patches = image.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size) # Flip one patch horizontally flip_idx = torch.randint(len(patches), (1,)).item() patches[flip_idx] = torch.flip(patches[flip_idx], dims=[-1]) question = f"Which patch is flipped horizontally?" answer = f"Patch number {flip_idx} is flipped." return question, answer @staticmethod def cropped_patch_inpainting(image, patch_size=32, mask_idx=None): patches = image.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size) if mask_idx is None: mask_idx = torch.randint(len(patches), (1,)).item() # Mask one patch masked_patches = patches.clone() masked_patches[mask_idx] = torch.zeros_like(masked_patches[mask_idx]) # Candidates: 4 random patches including correct one candidates = [patches[mask_idx]] + [patches[torch.randint(len(patches), (1,))] for _ in range(3)] correct_idx = 0 question = f"Which patch completes the missing region?" answer = f"The correct patch is option {correct_idx}." return question, answer, candidates @staticmethod def regional_depth_ordering(image, depth_map, num_regions=3): # Divide into regions, compute average depth h, w = image.shape[-2:] region_h, region_w = h // num_regions, w // num_regions regions = [] for i in range(num_regions): region_depth = depth_map[i*region_h:(i+1)*region_h, :].mean() regions.append((i, region_depth.item())) # Sort by depth sorted_regions = sorted(regions, key=lambda x: x[1]) question = "Rank these three image regions from closest to farthest from camera." answer = f"From closest to farthest: {[r[0] for r in sorted_regions]}" return question, answer @staticmethod def relative_position_prediction(image, depth_map): # Pick two points, ask about relative position pt1 = torch.tensor([image.shape[-2]//4, image.shape[-1]//4]) pt2 = torch.tensor([image.shape[-2]*3//4, image.shape[-1]*3//4]) depth1 = depth_map[pt1[0], pt1[1]] depth2 = depth_map[pt2[0], pt2[1]] question = f"Is point A closer or farther than point B?" answer = "closer" if depth1 < depth2 else "farther" return question, answer ``` **2. Generate Spatial-SSRL-81k Dataset** Create automatic dataset from existing image collections. ```python def generate_spatial_dataset(image_dir, depth_dir, output_size=81000): dataset = [] image_paths = list(Path(image_dir).glob("*.jpg"))[:output_size // 5] for img_path in image_paths: image = Image.open(img_path) depth_path = Path(depth_dir) / (img_path.stem + ".npy") depth = np.load(depth_path) if depth_path.exists() else None # Generate 5 QA pairs per image qa_pairs = [ SpatialPretext.shuffled_patch_reordering(image), SpatialPretext.flipped_patch_recognition(image), ] if depth is not None: qa_pairs.extend([ SpatialPretext.regional_depth_ordering(image, depth), SpatialPretext.relative_position_prediction(image, depth), ]) for q, a in qa_pairs: dataset.append({ 'image_path': str(img_path), 'question': q, 'answer': a }) return dataset[:output_size] ``` **3. Implement GRPO Training Loop** Train the model using Group Relative Policy Optimization. ```python class SpatialGRPOTrainer: def __init__(self, model, tokenizer): self.model = model self.tokenizer = tokenizer def compute_spatial_reward(self, generated_answer, reference_answer): # Exact match for spatial reasoning exact_match = float(generated_answer == reference_answer) # Format compliance (answer in proper boxes) has_box = '\\boxed{' in generated_answer format_reward = float(has_box) if has_box else 0.1 return 0.9 * exact_match + 0.1 * format_reward def train_step(self, batch, lr=1e-5): images, questions, answers = batch # Generate outputs inputs = self.tokenizer(questions, images) outputs = self.model.generate(inputs, max_length=256) generated_answers = self.tokenizer.decode(outputs) # Compute rewards per sample rewards = [ self.compute_spatial_reward(gen, ref) for gen, ref in zip(generated_answers, answers) ] # Group Relative Policy Optimization # Separate into groups by reward quartile sorted_indices = np.argsort(rewards) group_size = len(rewards) // 4 losses = [] for group_idx in range(4): group_start = group_idx * group_size group_end = (group_idx + 1) * group_size group_indices = sorted_indices[group_start:group_end] # Positive examples: high-reward samples # Negative examples: low-reward samples if group_idx >= 2: # High reward group loss = self.compute_grpo_loss( inputs[group_indices], answers[group_indices], positive=True ) else: # Low reward group loss = self.compute_grpo_loss( inputs[group_indices], answers[group_indices], positive=False ) losses.append(loss) total_loss = sum(losses) / len(losses) return total_loss.item() def compute_grpo_loss(self, inputs, targets, positive=True): outputs = self.model(**inputs) logits = outputs.logits # Standard language modeling loss target_ids = self.tokenizer(targets).input_ids loss = F.cross_entropy(logits, target_ids) # Scale by relative position return loss if not positive else loss * 0.5 ``` **4. Implement Cold-Start Supervised Fine-Tuning** Before RL, brief SFT stabilizes training. ```python def cold_start_sft(model, dataset, num_epochs=1): optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) for epoch in range(num_epochs): for batch in dataset: images, questions, answers = batch # Supervised next-token prediction outputs = model(images=images, questions=questions, answers=answers) loss = outputs.loss optimizer.zero_grad() loss.backward() optimizer.step() return model ``` ## Practical Guidance **When to Use**: - Improving spatial reasoning in existing VLMs - Robotics or navigation tasks requiring 3D understanding - Models struggling with geometric reasoning despite strong language performance **Hyperparameters**: - num_epochs_sft: 1 (brief cold-start is sufficient) - lr_grpo: 1e-5 (lower than supervised learning) - reward_accuracy_weight: 0.9 (emphasize correctness over format) **When NOT to Use**: - Models already trained on spatial-intensive datasets - High-precision vision tasks requiring pixel-level accuracy - Scenarios where additional training is prohibitively expensive **Pitfalls**: - **Insufficient pretext task diversity**: Five tasks cover basic geometry; add task-specific variants for your domain - **Depth map quality**: Depth-dependent tasks require accurate depth estimates; synthetic or coarse depth can mislead - **Overfitting to pretext tasks**: Spatial reasoning sometimes doesn't transfer to downstream tasks; validate transfer explicitly **Key Insight**: Unlike typical self-supervised learning, you don't need paired data or external annotation. The geometric properties of images themselves provide infinite supervision. ## Reference arXiv: https://arxiv.org/abs/2510.27606