K-PWM: Control-Oriented Structured World Models under Partial Observation

World models enable prediction, planning, and decision-making by learning internal simulators of environment dynamics. Although substantial progress has been made in world models from pixel-based observations, many physical systems are instead observed through vector-based measurements that are noisy and only partially informative of the underlying state. We introduce K-PWM, a Koopman-structured probabilistic world model for prediction and control under partial observation. K-PWM combines a parameterized Koopman-structured state-space model with a nonlinear observation decoder, and performs probabilistic latent-state inference directly through the learned model parameters rather than through a separately trained inference network. We train K-PWM using a generalized expectation conditional maximization (GECM) procedure with principled initialization. Across Gymnasium MuJoCo tasks, K-PWM improves sample efficiency, yields reliable long-horizon predictions, and supports downstream control through model predictive path integral (MPPI) control. These results suggest that structured probabilistic latent dynamics provide a useful route toward data-efficient world modeling and control in partially observable vector-valued settings.