Efficient Gaussian Processes for Model-based Planning

Report | Slides | GitHub repository

Summary

How can a planning agent learn useful dynamics from limited interaction without making online inference prohibitively expensive? This was the central question of my final-year study project at NUS, supervised by Prof. Jonathan Scarlett. The work was exploratory: I reproduced and combined ideas from scalable Gaussian-process inference and TD-MPC-style planning, then evaluated several small modifications rather than presenting a finished new algorithm.

Model-based reinforcement learning (MBRL) learns a dynamics model and uses it to evaluate candidate actions before execution. Neural dynamics models scale well but can be data-hungry and biased. Gaussian processes (GPs) provide flexible non-parametric predictions and calibrated uncertainty in low-data regimes, but exact training and inference scale poorly. The project asked where scalable GP approximations could improve a TD-MPC-style planner without losing its practical runtime.

Key Ideas and Tricks

1. TD-MPC-style planning. Short latent rollouts reduce the number of expensive model evaluations required during trajectory optimization.

2. GP correction of MLP dynamics. An MLP provides a scalable base prediction, while independent GPs learn residual corrections that can, in principle, be evaluated in parallel.

3. Decoupled training and inference. Hyperparameters are optimized offline; online updates use variational conditioning so that prediction cost depends on the inducing set rather than the full dataset.

4. Deep kernel learning (DKL). A learned feature map can reduce the burden on the kernel in high-dimensional state-action spaces, at the cost of additional optimization and weaker sample efficiency in some experiments.

Some Findings

1. GP-corrected models often matched or outperformed the TD-MPC baseline across the tested continuous-control tasks, with a modest and largely parallelizable runtime overhead.

2. Matérn kernels were often more effective in environments with abrupt, contact-driven dynamics. A plausible explanation is that their weaker smoothness assumptions better match these transitions; I regard this as an empirical observation rather than a definitive causal conclusion.

3. DKL generally required fewer inducing points to reach a target approximation error, but sometimes introduced extra runtime and reduced sample efficiency.

4. Higher-dimensional environments still required many inducing points. I explored local kernel interpolation and a preliminary “dynamical local projection” method, but neither was mature enough to resolve this limitation.

Reflections

By the end of the project, Gaussian processes had lost some of their mystique. They are not a universally superior alternative to neural networks; they offer a different collection of inductive biases and computational tradeoffs. The experiments suggest that hybrid models are most useful when uncertainty and sample efficiency matter, the effective dimension remains manageable, and the implementation can exploit parallelism. Modern numerical methods make these classical models worth revisiting, but they do not remove the need to choose them judiciously.