Threading the Needle: Targeted Keyframe Sampling for 4D World Models in Precise Robotic Manipulation

Abstract

High-precision robotic tasks, such as peg insertion, pose a challenge for learned world models. The critical states that determine success occur within an extremely narrow portion of each task trajectory, causing uniform sampling to underrepresent them. This thesis investigates whether targeted data selection improves the generation quality and downstream success of a 4D latent world model on the ManiSkill3 PegInsertionSide-v1 benchmark. We propose \textit{Pre-insertion Keyframe Sampling}, an algorithm that identifies and injects a single geometrically aligned pre-insertion frame into training data, and introduce a \textit{Peg-Hole Alignment Metric} to evaluate 3D generation quality beyond standard image metrics. Our results demonstrate that while the intervention yields statistically significant improvements in translational alignment, these gains do not translate to higher task success rates. A ground-truth subgoal ablation reveals a 56-percentage-point performance gap, identifying the structural fidelity of 3D generations as the primary bottleneck. This suggests that while data coverage improves spatial precision, high-precision control requires increasing the geometric resolution and structural integrity of the world model's output.