Yihao Sun

Ph.D. Student in Computer Science at Mila - Quebec AI Institute & Université de Montréal. Advised by Prof. Pierre-Luc Bacon.

Before starting my Ph.D., I completed my M.Sc. in the LAMDA Group at Nanjing University.

My research centers on reinforcement learning for general-purpose embodied intelligence, with a particular focus on world models and model-based reinforcement learning. I am interested in building generalist world models for embodied environments and developing algorithms that use them for efficient policy learning, planning, and adaptation.

Email GitHub Google Scholar CV

Selected Publications

*: equal contribution, see Google Scholar for the full list

Placeholder illustration for the VLA-MBPO research paper

Towards Practical World Model-based Reinforcement Learning for Vision-Language-Action Models

Zhilong Zhang^*, Haoxiang Ren^*, Yihao Sun^*, Yifei Sheng, Haonan Wang, Haoxin Lin, Zhichao Wu, Pierre-Luc Bacon, Yang Yu

ICML, 2026

paper project page

VLA-MBPO integrates a series of pratical techniques to train VLAs in world models, including (i) adapting unified multimodal models (UMMs) for data-efficient world modeling; (ii) an interleaved view decoding mechanism to enforce multi-view consistency; and (iii) chunk-level branched rollout to mitigate error compounding.

ADM-v2: Pursuing Full-Horizon Roll-out in Dynamics Models for Offline Policy Learning and Evaluation

Haoxin Lin, Siyuan Xiao, Yi-Chen Li, Zhilong Zhang, Yihao Sun, Chengxing Jia, Yang Yu

ICLR, 2026

paper code

ADM-v2 not only enhances the accuracy of direct predictions for making fullhorizon rollouts but also supports parallel estimation of the any-step prediction uncertainty to improve efficiency.

Any-step Dynamics Model Improves Future Predictions for Online and Offline Reinforcement Learning

Haoxin Lin, Yu-Yan Xu, Yihao Sun, Zhilong Zhang, Yi-chen Li, Chengxing Jia, Yang Yu

ICLR, 2025

paper code

In this paper, we propose the Any-step Dynamics Model (ADM) for MBRL to mitigate the compounding error by reducing bootstrapping prediction to direct prediction.We design two algorithms, We then build ADMPO-ON and ADMPO-OFF, which apply ADM in online and offline model-based frameworks, respectively.

Flow to Better: Offline Preference-based Reinforcement Learning via Preferred Trajectory Generation

Zhilong Zhang^*, Yihao Sun^*, Junyin Ye, Tian-shuo Liu, Jiaji Zhang, Yang Yu

ICLR, 2024

paper project page

FTB considers the offline preference-based RL as a generative modeling problem, which leverages the pairwise preference relationship to guide a generative model in producing preferred trajectories.

Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning

Yihao Sun^*, Jiaji Zhang^*, Chengxing Jia, Haoxin Lin, Junyin Ye, Yang Yu

ICML, 2023

paper code

MOBILE proposes a novel uncertainty quantification for uncertainty-driven model-based offline RL, which can better estimate the ideal uncertainty quantification, i.e., the Bellman error

Open Source Projects

VLARLKit

github repository

A researcher-friendly PyTorch reinforcement-learning library for Vision-Language-Action models, with clean modular design and decoupled env denpendency.

OfflineRL-Kit

github repository

An elegant PyTorch offline-reinforcement-learning library with clear, extensible implementations of model-free and model-based algorithms.