CV

Education

Publications

* indicates equal contribution. † indicates co-advising. First-author or co-first-author works are highlighted.

WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems

WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems

Yuchen Wang*, Jiangtao Kong*, Sizhe Wei, Xiaochang Li, Haohong Lin, Hongjue Zhao, Tianyi Zhou, Lu Gan, and Huajie Shao.

International Conference on Machine Learning (ICML), 2026

Spotlight (2.2%)

Trajectory world models play a crucial role in robotic dynamics learning, planning, and control. While recent works have explored trajectory world models for diverse robotic systems, they struggle to scale to a large number of distinct system dynamics and overlook domain knowledge of physical structures. To address these limitations, we introduce WestWorld, a knowledge-encoded scalable trajectory world model for diverse robotic systems. To tackle the scalability challenge, we propose a novel system-aware Mixture-of-Experts (Sys-MoE) that dynamically combines and routes specialized experts for different robotic systems via a learnable system embedding. To further enhance zero-shot generalization, we incorporate domain knowledge of robot physical structures by introducing a structural embedding that aligns trajectory representations with morphological information. After pretraining on 89 complex environments spanning diverse morphologies across both simulation and real-world settings, WestWorld achieves significant improvements over competitive baselines in zero- and few-shot trajectory prediction. Additionally, it shows strong scalability across a wide range of robotic environments and significantly improves performance on downstream model-based control for different robots. Finally, we deploy our model on a real-world Unitree Go1, where it demonstrates stable locomotion performance. The code is available at https://github.com/511205787/WestWorld.

A Generalizable Physics-guided Causal Model for Trajectory Prediction in Autonomous Driving

A Generalizable Physics-guided Causal Model for Trajectory Prediction in Autonomous Driving

Zhenyu Zong, Yuchen Wang, Haohong Lin, Lu Gan, and Huajie Shao.

2026 IEEE International Conference on Robotics and Automation (ICRA);, 2026

Trajectory prediction for traffic agents is critical for safe autonomous driving. However, achieving effective zero-shot generalization in previously unseen domains remains a significant challenge. Motivated by the consistent nature of kinematics across diverse domains, we aim to incorporate domain-invariant knowledge to enhance zero-shot trajectory prediction capabilities. The key challenges include: 1) effectively extracting domain-invariant scene representations, and 2) integrating invariant features with kinematic models to enable generalized predictions. To address these challenges, we propose a novel generalizable Physics-guided Causal Model (PCM), which comprises two core components: a Disentangled Scene Encoder, which adopts intervention-based disentanglement to extract domain-invariant features from scenes, and a CausalODE Decoder, which employs a causal attention mechanism to effectively integrate kinematic models with meaningful contextual information. Extensive experiments on real-world autonomous driving datasets demonstrate our method’s superior zero-shot generalization performance in unseen cities, significantly outperforming competitive baselines. The source code is released at https://github.com/ZY-Zong/Physics-guided-Causal-Model.

A Generalizable Physics-Enhanced State Space Model for Long-Term Dynamics Forecasting in Complex Environments

A Generalizable Physics-Enhanced State Space Model for Long-Term Dynamics Forecasting in Complex Environments

Yuchen Wang, Hongjue Zhao, Haohong Lin, Enze Xu, Lifang He, and Huajie Shao.

Proceedings of the 42nd International Conference on Machine Learning (ICML), 2025

This work aims to address the problem of long-term dynamic forecasting in complex environments where data are noisy and irregularly sampled. While recent studies have introduced some methods to improve prediction performance, these approaches still face a significant challenge in handling long-term extrapolation tasks under such complex scenarios. To overcome this challenge, we propose Phy-SSM, a generalizable method that integrates partial physics knowledge into state space models (SSMs) for long-term dynamics forecasting in complex environments. Our motivation is that SSMs can effectively capture long-range dependencies in sequential data and model continuous dynamical systems, while the incorporation of physics knowledge improves generalization ability. The key challenge lies in how to seamlessly incorporate partially known physics into SSMs. To achieve this, we decompose partially known system dynamics into known and unknown state matrices, which are integrated into a Phy-SSM unit. To further enhance long-term prediction performance, we introduce a physics state regularization term to make the estimated latent states align with system dynamics. Besides, we theoretically analyze the uniqueness of the solutions for our method. Extensive experiments on three real-world applications, including vehicle motion prediction, drone state prediction, and COVID-19 epidemiology forecasting, demonstrate the superior performance of Phy-SSM over the baselines in both long-term interpolation and extrapolation tasks. The code is available at https://github.com/511205787/Phy_SSM-ICML2025.

Accelerating Neural ODEs: A Variational Formulation-based Approach

Accelerating Neural ODEs: A Variational Formulation-based Approach

Hongjue Zhao, Yuchen Wang, Hairong Qi, Zijie Huang, Han Zhao, Lui Sha, and Huajie Shao.

The Thirteenth International Conference on Learning Representations (ICLR), 2025

Neural Ordinary Differential Equations (Neural ODEs or NODEs) excel at modeling continuous dynamical systems from observational data, especially when the data is irregularly sampled. However, existing training methods predominantly rely on numerical ODE solvers, which are time-consuming and prone to accumulating numerical errors over time due to autoregression. In this work, we propose VF-NODE, a novel approach based on the variational formulation (VF) to accelerate the training of NODEs. Unlike existing training methods, the proposed VF-NODE implements a series of global integrals, thus evaluating Deep Neural Network (DNN)-based vector fields only at specific observed data points. This strategy drastically reduces the number of function evaluations (NFEs). Moreover, our method eliminates the use of autoregression, thereby reducing error accumulations for modeling dynamical systems. Nevertheless, the VF loss introduces oscillatory terms into the integrals when using the Fourier basis. We incorporate Filon's method to address this issue. To further enhance the performance for noisy and incomplete data, we employ the natural cubic spline regression to estimate a closed-form approximation. We provide a fundamental analysis of how our approach minimizes computational costs. Extensive experiments demonstrate that our approach accelerates NODE training by 10 to 1000 times compared to existing NODE-based methods, while achieving higher or comparable accuracy in dynamical systems. The code is available at https://github.com/ZhaoHongjue/VF-NODE-ICLR2025.

A Deep Transfer Operator Learning Method for Temperature Field Reconstruction in a Lithium-Ion Battery Pack

A Deep Transfer Operator Learning Method for Temperature Field Reconstruction in a Lithium-Ion Battery Pack

Yuchen Wang, Can Xiong, Changjiang Ju, Genke Yang, Yu-wang Chen, and Xiaotian Yu.

IEEE Transactions on Industrial Informatics (IF=11.7), 2024

Nonuniform thermal behavior in lithium-ion battery packs can accelerate aging, leading to inconsistent cell performance. If not adequately monitored and managed, this heating can give rise to unwanted side reactions, fires, and explosions, underscoring the criticality of temperature field reconstruction. In recent years, data-driven methods have gained popularity for addressing the temperature field reconstruction problem. However, many existing data-driven approaches require retraining when system parameters change, such as the initial temperature distribution or working conditions. This article presents a deep transfer operator learning method named physics-informed adversarial networks. The model architecture incorporates transformer blocks to capture comprehensive time and space features. Additionally, to enhance interpretability and generalization, the model introduces two effective mechanisms: 1) the integration of thermal partial differential equations to ensure compliance with physical laws; and 2) the application of domain adversarial mechanism in transfer learning to extract domain-invariant feature representations. These mechanisms enable the model to effectively reconstruct the temperature field, even in unencountered scenarios during training. The proposed method is validated under real-world energy storage working conditions, demonstrating superior performance compared to state-of-the-art deep learning methods. Notably, the approach exhibits excellent performance even when confronted with the limited availability of training data.

Accelerating Neural Differential Equations for Irregularly-Sampled Dynamical Systems Using Variational Formulation

Accelerating Neural Differential Equations for Irregularly-Sampled Dynamical Systems Using Variational Formulation

Hongjue Zhao, Yuchen Wang, Hairong Qi, Jiajia Li, Lui Sha, Han Zhao, and Huajie Shao.

ICLR 2024 Workshop on AI4DifferentialEquations In Science, 2024

Neural ODEs have exhibited remarkable capabilities in continuously modeling dynamical systems from observational data. However, existing training methods, often based on adaptive-step-size numerical ODE solvers, are time-consuming and may introduce additional errors. Despite recent attempts to address these issues, existing methods still rely heavily on numerical ODE solvers and lack efficient solutions. In this work, we propose the Fast-VF Neural ODE, a novel approach based on variational formulation (VF) to accelerate the training of Neural ODEs for dynamical systems. To further mitigate the influence of oscillatory terms in the VF loss, we incorporate Filon's method into our design. Extensive experimental results show that our method can accelerate the training of Neural ODEs by 10× to 100× compared to baseline methods while achieving comparable accuracy in irregularly sampled dynamical systems.

Temperature state prediction for lithium-ion batteries based on improved physics informed neural networks

Temperature state prediction for lithium-ion batteries based on improved physics informed neural networks

Yuchen Wang, Can Xiong, Yiming Wang, Po Xu, Changjiang Ju, Jianghao Shi, Genke Yang, and Jian Chu.

Journal of Energy Storage (IF=8.9), 2023

Heat generation significantly influences the performance of lithium-ion batteries and also hinders their application. Precise prediction of battery temperature can provide feedback for system monitoring and enable safe and efficient battery operation. However, battery temperature prediction remains extremely challenging because irreversible heat increases as batteries age across the life cycle. To address this problem, we propose a battery informed neural network (BINN). The framework incorporates battery physical models into long short-term memory (LSTM)-based networks in an end-to-end manner for battery temperature prediction. A multi-head attention mechanism is introduced to capture information from longer time series. Physical parameters of the battery electrical model, heat generation model, and thermal model are automatically learned during training, and the irreversible heat changes caused by aging are represented through these physical parameters. Temperature prediction over the full life cycle of lithium-ion batteries is evaluated under different working conditions. The results show that BINN is interpretable and achieves better generalization and transferability than traditional learning-based methods.

Honors & Awards

Services