Technical Summary: PLUME

Problem Statement

Dexterous multi-finger manipulation is highly sensitive to physical parameters such as object shapes, poses, and friction coefficients. These parameters are often uncertain at deployment, leading to unknown system dynamics. While simulation allows for large-scale data collection with known parameters, policies trained via standard domain randomization often fail to generalize to precise tasks (e.g., screwdriver turning) because optimal manipulation strategies may need to change based on specific parameter values. Furthermore, collecting high-quality data for complex multi-finger tasks is challenging, often resulting in heterogeneous datasets containing sub-optimal trajectories or failures. Existing world modeling methods typically rely on fixed-horizon histories to implicitly account for dynamics variations, which is insufficient for precise tasks sensitive to such variations. Additionally, offline reinforcement learning methods generally do not focus on the joint problem of system identification and policy learning, while online system identification methods often require complex reward engineering and do not support planning.

Methodology

The authors propose PLUME (Probabilistic Latent Unified world Modeling and parameter Estimation), a framework that jointly learns a world model and estimates a belief over latent physical parameters.

Core Architecture

PLUME consists of three main components trained end-to-end:

1. Parameter Representation Learning: A Wasserstein Auto-Encoder (WAE) with a Transformer backbone learns a unified latent representation ($z_t$) of heterogeneous variables, including physical parameters (e.g., friction, shape) and rewards. The encoder tokenizes these inputs and uses cross-attention pooling to produce a latent embedding. A standard normal prior is imposed to facilitate reasonable initialization of parameter beliefs.
2. History Modeling: A Mamba-2 encoder compresses the trajectory history (past actions and observations) into a context vector ($c_t$), retaining useful historical context that simple Markov assumptions might lose, especially with noisy observations or limited particles.
3. Joint Flow Matching: A flow matching model serves as the unified world model. It learns to estimate latent parameters and model system dynamics conditioned on those parameters. The model uses a transformer backbone and an auxiliary mask to handle different conditional distributions:
   • Parameter Estimation: $p(Z_{t+1} | Z_t, y_t, u_{t+1}, y_{t+1}, c_t)$
   • Rollout Generation: $p(u_{t+1}, y_{t+1}, Z_{t+1} | Z_t, y_t, c_t)$
   • Inverse Dynamics: $p(u_{t+1} | Z_t, y_t, Z_{t+1}, y_{t+1}, c_t)$

Training and Inference

• Training: The model is trained on offline datasets containing trajectories with varying parameter values. The loss function combines a reconstruction loss for parameters, a Maximum Mean Discrepancy (MMD) regularization to align the latent space with a standard normal, and the flow matching loss. Crucially, gradients from the flow matching loss are not backpropagated to the parameter encoder to prevent the model from optimizing its own target at the expense of representational capacity.
• Planning (Deployment): At deployment, PLUME maintains a belief over latent parameters represented as a set of particles ($Z_t$). It employs a receding-horizon trajectory sampling approach:
  1. Rollout: The model autoregressively samples multiple trajectories conditioned on the current observation, history, and particle belief.
  2. Scoring: Trajectories are scored based on a combination of decoded rewards and estimated transition likelihoods. The likelihood is estimated using an inverse dynamics model; low likelihoods penalize high-reward trajectories that exhibit unrealistic dynamics (Out-of-Distribution detection).
  3. Execution & Update: The first action of the highest-scoring trajectory is executed. The latent belief is then updated online using the new observation, allowing the world model to adapt to the true dynamics without re-training or fine-tuning.

Key Contributions

1. Unified Flow-Matching World Model: A model that estimates system parameters to condition world model rollouts, enabling a single policy to adapt to varying dynamics.
2. End-to-End Latent Representation: A method to learn a unified latent space for qualitatively different parameters (friction, shape) and rewards, optimized jointly with the world model to be useful for dynamics learning.
3. Likelihood-Weighted Planning: A planning strategy that utilizes world model rollouts and estimated transition likelihoods to optimize robot actions, effectively filtering out unrealistic high-reward plans.

Experimental Results

The authors evaluated PLUME on four simulated tasks (screwdriver turning, valve turning, disk flicking, bucket lifting) and one hardware task (screwdriver turning).

• Baselines: Comparisons were made against Unified World Model (UWM), Flow Q-Learning (FQL), and Decision Diffuser (DD).
• Performance:
  • Simulated Screwdriver: PLUME achieved an 81.3% validity rate, significantly outperforming UWM (70.7%), FQL (21.3%), and DD (19.3%). It also reduced the goal distance by 31% compared to the best baseline.
  • Hardware Screwdriver: PLUME achieved a 93.3% validity rate (28/30 trials) with zero-shot transfer from simulation, outperforming UWM (86.7%).
  • Other Tasks: PLUME significantly outperformed all baselines in valve turning and disk flicking tasks (measured by goal distance) and was the only method to consistently succeed in the DexArt bucket lifting task.
• Ablations: The paper notes that PLUME's ability to reason about partially observable state variables (position, velocity, orientation) in addition to physical parameters contributes to its success, as evidenced by accurate reward estimation.

Significance and Claims

The paper claims that for precise and sensitive multi-finger manipulation tasks, explicitly estimating and conditioning on representations of system parameters leads to superior performance compared to methods that learn parameter-agnostic policies or model variation implicitly.

Key claims include:

• Zero-Shot Transfer: PLUME can deploy a simulation-trained policy to hardware tasks (screwdriver turning) without fine-tuning, adapting to the true dynamics through online parameter inference.
• Robustness to Data Quality: The method can learn useful models from heterogeneous datasets containing sub-optimal trajectories and failures, leveraging reward labels to use failure data for dynamics learning while avoiding problematic actions through planning.
• Adaptability: Unlike methods requiring re-training or fine-tuning, PLUME aligns its world model to true dynamics at deployment by updating its belief over parameters, allowing a single model to handle diverse physical conditions.

The authors acknowledge limitations, specifically the need for ground-truth parameter values during training (available in simulation but not always in the real world) and the quadratic memory scaling of transformers with respect to the number of particles. They suggest future work could involve fine-tuning pre-trained models on real-world data to align priors without requiring ground-truth labels.