Ziang Fu

Learning representations that make good action sequences easier to find.

Accurate prediction alone does not ensure that a world model is easy to plan through. Control-Geometry Straightening (CGS) aligns the geometry of actions with corresponding latent changes, using a single auxiliary loss on pixel–action pairs.

Across four control environments, CGS improves planning with fewer sampled candidates and refinement steps. Our analysis connects this geometry to more balanced terminal-cost curvature over the planning horizon.