GPU-Accelerated Astrodynamics World Models for Spacecraft Rendezvous and Proximity Operations

Duncan Eddy; Isaac R. Ward; Grace Ra Kim; Mykel J. Kochenderfer

Source-linked AI summary

GPU-Accelerated Astrodynamics World Models for Spacecraft Rendezvous and Proximity Operations

Duncan Eddy, Isaac R. Ward, Grace Ra Kim, Mykel J. Kochenderfer

arXiv:2609.03067v1Sep 2, 2026cs.ROeess.SY

TL;DR

Autonomous spacecraft rendezvous must handle uncertain dynamics and sensing, but world models had not previously been applied to space operations. This paper introduces a GPU-enabled ISS docking simulator and OWM, a multimodal probabilistic world model, and reports improved docking, generalization, and anomaly-detection results over supported baselines.

Problem
Autonomous rendezvous must reason over uncertain dynamics and sensing while performing precise, safe docking maneuvers.
Method
The paper combines AstroJAX GPU simulation with OWM, a transformer that fuses kinematic and visual observations and predicts future-observation distributions.
Results
53% versus 29% docking success across eight ports, 40% versus 17% on held-out ports, and 98% anomaly-classification accuracy.
Takeaways & Limitations
The approach demonstrates world-model-based ISS docking with improved performance on unseen ports and runtime anomaly detection.
Takeaways & Limitations
The reinforcement-learning baseline required adaptation from 20 Hz to 1 Hz control because the terminal docking reward was effectively discounted away.

Abstract

from arXiv · show

World models are an emerging paradigm in representation learning in which an agent jointly learns state-action dynamics and observation models from offline trajectory data, enabling multi-step planning and trajectory prediction with uncertainty estimates. They have shown strong results in robotics and game environments, but, to the best of our knowledge, have not previously been applied to the space domain. This paper introduces a world model-based approach to cooperative and non-cooperative spacecraft rendezvous and proximity operations. First, we introduce an open-source, JAX-based International Space Station (ISS) docking environment supporting parallel GPU simulation of spacecraft orbit and attitude dynamics, generating the thousands of state-action transitions that world model training requires. Second, we introduce Out-of-this-World-Model, a transformer-based world model that encodes relative kinematic states and body-fixed camera imagery into a latent state and predicts its evolution under commanded thrusts and torques using one-step flow matching. It produces a distribution over future observations, capturing stochastic dynamics and per-timestep uncertainty, and outperforms DreamerV3-style posterior-correction baselines with fewer trainable parameters and hyperparameters. Third, we apply the approach to a capsule autonomously docking with the ISS under keep-out-zone constraints, demonstrating improved sample efficiency and task performance over reinforcement learning baselines (53% versus 29% docking success across ports), better out-of-distribution generalization (on held-out ports the world model more than doubles baseline success, 40% versus 17%), and detection of anomalous objects encountered during approach with 98% classification accuracy. We open-source the simulation environment and model architecture to enable further study of this paradigm.

1. INTRODUCTION

Spacecraft rendezvous is increasingly important but remains risky because autonomous systems must reason over uncertain dynamics and sensing before executing precise, constrained maneuvers. This paper introduces a GPU-accelerated simulation framework and a world model that predicts uncertain future observations for ISS docking, improving performance and anomaly detection.

Rendezvous and docking are becoming routine across civil, military, commercial, and human-spaceflight operations, including every crew and cargo visit to the ISS.
Centimeter-level docking precision and collision risk make autonomous rendezvous safety-critical, requiring prediction of control consequences and early recognition of off-nominal scenarios.
Autonomous rendezvous must fuse uncertain measurements from cooperative and non-cooperative sensors, with cameras supporting communication-independent operation but providing less accurate relative-state estimates.
Existing learning methods struggle with long-horizon prediction or model only dynamics, while world models jointly learn state-action dynamics and observation models to predict future observations.
AstroJAX provides open-source, massively parallel GPU simulation of orbit and attitude dynamics, with an ISS-docking environment for generating world-model training data at scale.
OWM fuses kinematic and visual observations into a shared latent space and predicts distributions over future observations using a one-step flow-matching head.
53% versus 29% docking success across eight ports, with 40% versus 17% on held-out berthing ports, while predictive uncertainty classified anomalous approach sequences 98% of the time.
The authors open-source the simulation environments, model architecture, and training datasets to support further study.

2. BACKGROUND

The background connects world models with spacecraft GNC, emphasizing their ability to learn predictive observation models and uncertainty from data rather than relying solely on analytic decompositions. Prior work spans learned latent dynamics, classical control, reinforcement learning, and specialized simulation infrastructure, but does not apply world models to space operations.

World Models: World models learn compressed representations and predict future system behavior, with Dreamer-style methods combining latent priors and observation-conditioned posterior correction.
World Models: JEPA-style approaches predict future latent representations without reconstructing observations, reducing cost but risking encoder collapse to a constant representation.
World Models: Flow matching transports noise samples toward a data distribution, while shortcut models enable prediction with one large integration step instead of many small ones.
World Models: Prior learned-system-model work covers fluid simulation and robotic manipulation, where predictive uncertainty has been used to detect failures during real-world tasks.
Rendezvous and Proximity Operations: Reinforcement learning is the most studied learning-based approach to rendezvous and proximity operations, including policy optimization, meta-learning, image-based navigation, and cislunar applications.
Rendezvous and Proximity Operations: Simulation infrastructure such as BSK-RL and Basilisk supports learning autonomous inspection and small-body science operations.
Rendezvous and Proximity Operations: Classical docking methods handle thrust, approach-cone, soft-docking, and motion constraints, but depend on the models and processed relative-state estimates they are given.
Rendezvous and Proximity Operations: Neural spacecraft-dynamics methods model dynamics alone and face long-horizon accuracy or training challenges, whereas world models jointly learn dynamics and observation prediction.

3. METHODS

The paper combines a GPU-accelerated JAX simulation stack with a transformer world model that predicts uncertain spacecraft observations and latent dynamics from recorded trajectories. The model integrates kinematic and visual inputs, rolls forward autoregressively in latent space, and is trained with flow matching and observation-based reconstruction losses.

Problem formulation: The environment models docking as a partially observable decision process in which actions advance spacecraft dynamics while noisy kinematic measurements and camera images form observations.The world model predicts the next observation distribution from a history of past observations and actions rather than directly observing the underlying state.
World model architecture: The model encodes kinematic states and camera frames into per-timestep latent tokens processed by factorized spatial and causal temporal transformer attention.Spatial attention fuses modalities within each timestep, while temporal attention models history across timesteps; factorization reduces full-sequence attention cost.
World model architecture: Flow matching samples latent residuals that update the current latent state, enabling autoregressive latent-space rollouts and distributions over future observations with uncertainty.Different Gaussian noise realizations produce different futures, while the decoder maps predicted latents to images and states.
Training: Training combines rectified-flow prediction of latent residuals with per-modality decoding and latent round-trip reconstruction losses.The reconstruction terms keep the shared latent representation grounded in camera and kinematic observations while limiting decoder computation through sampled rollout steps.
Simulation and data generation: AstroJAX and the ISS-docking environment generate high-throughput training trajectories using parallel, GPU-accelerated JAX simulation.The framework implements astrodynamics and environment logic as JAX functions, allowing parallel CPU/GPU execution and fused trajectory rollouts.

4. EXPERIMENTS

The experiments generate diverse noisy ISS-docking trajectories, train and evaluate world-model planning against PPO, and assess prediction quality, docking, generalization, and anomaly-related capabilities.

4.2 Evaluation and Planning: The study compares world-model MPPI planning with PPO across prediction quality, in-distribution docking, unseen-port generalization, and anomaly detection.MPPI rolls candidate action sequences through the learned model and replans after executing short action chunks.
4.1 Data Generation: The ISS environment generates trajectories from random, orbit, and docking policies, covering undirected motion, station circumnavigation, and approaches to sampled ports.Episodes can end in docking, collision, departure, or horizon timeout.
4.1 Data Generation: Training uses five of eight docking ports, while Harmony zenith, Poisk zenith, and Unity nadir are held out for generalization evaluation.The held-out ports have unseen goal poses, approach corridors, and visual context.
4.1 Data Generation: The dataset records true states, noisy observations, actions, rewards, and first-person camera views under cooperative and non-cooperative sensor-noise models.Behavior policies act on noisy measurements, preserving aleatoric uncertainty from acting on observed rather than true states.
4.3 Baseline Algorithms: PPO required substantial adaptation, including 1 Hz control, a soft keep-out zone, and reward shaping, because the native 20 Hz setup made terminal docking rewards effectively unreachable.At γ = 0.995, the 7200-step discount is approximately 2×10^-16.
4.3 Baseline Algorithms: The baseline reward combines distance, attitude, velocity, body-rate, actuator, collision, docking, and departure terms using normalized pseudo-Huber losses and proximity-based bonuses.The far field maintains a constant pull toward the port while the near field remains smooth; rotational terms are weighted by range.

5. RESULTS

The world-model planner achieves lower prediction error than a Dreamer-style baseline and improves docking acquisition and generalization, while predictive uncertainty supports anomaly detection. Its performance depends on how success and safety are defined, with stronger results for acquisition than strict terminal contact.

5.1 Training and Prediction Quality: Both world-model variants converge within 16 epochs, while flow matching achieves lower state and image prediction errors than the Dreamer-style baseline with fewer trainable parameters.Validation uses fully autoregressive deployment conditions rather than teacher forcing.
5.1 Training and Prediction Quality: Non-cooperative sensor noise degrades long-horizon state prediction more than image prediction, while image quality remains nearly insensitive because the camera channel is noise-free.The state degradation is attributed to perturbations applied directly to state observations.
5.2 Docking Performance: 61% versus 36%: collision-voided success at 1 m on training ports favors the world-model planner over PPO, which the planner also surpasses at several individual ports.The planner reaches Harmony forward at 76% and Zvezda aft at 70%, whereas PPO never reaches those ports.
5.2 Docking Performance: 48%, 42%, and 30%: the planner succeeds at all three held-out ports, whereas PPO reaches only Unity nadir at 50% and fails at both zenith ports.The held-out ports were absent from training trajectories, so the comparison tests out-of-distribution generalization.
5.2 Docking Performance: 36% versus 13%: under full contact conditions at 1 m on training ports, PPO exceeds the planner, while position-only success favors the planner across tolerances.Full contact conditions additionally constrain velocity, attitude error, and body rate; the paper attributes much of the gap to tuning effort.
5.2 Docking Performance: 0% versus 88%: the planner never escapes the operational domain across eight ports, while PPO escapes at each zenith port; the planner consequently attempts a full approach on every rollout.The planner’s higher collision rate partly reflects making approaches that baseline rollouts abandon before close approach.
5.3 Anomaly Detection: 98%: predictive uncertainty correctly classifies approach sequences with an anomalous docked vehicle versus those without one.The model had not seen a second visiting vehicle during training, and prediction-observation disagreement concentrates on the anomalous object.

6. CONCLUSIONS

The paper presents AstroJAX and OWM for GPU-accelerated simulation and world-model-based ISS docking, combining learned prediction with sampling-based planning and uncertainty monitoring. Results show transfer to unseen ports and runtime anomaly classification, while cost tuning and safety integration remain future work.

6. CONCLUSIONS: AstroJAX provides open-source, massively parallel GPU simulation of orbit and attitude dynamics for generating world-model training data at scale.The paper reports a 500,000-transition training corpus.
6. CONCLUSIONS: OWM fuses camera imagery and kinematic measurements into a shared latent space and predicts distributions over future observations with a one-step flow-matching head.It composes with model predictive path integral control for ISS docking.
6. CONCLUSIONS: 40% versus 17%: on held-out ports, the world-model planner more than doubles reinforcement-learning baseline docking success, while predictive uncertainty classifies anomalous-vehicle approaches with 98% accuracy.The planner also reaches ports the baselines fail to acquire and never leaves the operational domain.
6. CONCLUSIONS: Future work focuses on tuning MPPI terminal attitude and body-rate costs, using anomaly uncertainty for collision mitigation, and integrating world-model planners with safety-aware planning.The authors expect improved cost weighting to reduce collisions and narrow the gap under full contact conditions.