Technical Summary: Generalised Linear Models in Deep Bayesian RL with Learnable Basis Functions (GLiBRL)

Problem Statement

Bayesian Reinforcement Learning (BRL) offers a principled framework for generalization in Meta-Reinforcement Learning (Meta-RL) by explicitly modeling uncertainty over transition and reward dynamics. However, classical BRL methods are limited by their reliance on known functional forms for these models, restricting their applicability to diverse tasks. Recent deep BRL approaches attempt to learn these forms using neural networks but face significant theoretical and practical hurdles:

1. Intractable Inference: Directly applying neural networks to the joint space of data and task parameters renders exact Bayesian inference intractable, forcing reliance on Variational Inference (VI).
2. VI Limitations: VI introduces challenges such as high-variance Monte Carlo estimates, amortization gaps, and posterior collapse, often resulting in indistinct task representations that degrade policy performance.
3. Permutation Variance: Many deep BRL methods use sequential models (e.g., RNNs, Transformers) to summarize task histories. These are permutation-variant, making them incompatible with sample-efficient off-policy algorithms like Soft Actor-Critic (SAC) without severe performance degradation.
4. Noise Assumptions: Existing methods often assume known noise distributions or fix neural network weights at test-time (effectively performing Maximum Likelihood Estimation rather than Bayesian inference), leading to predictive errors in unseen tasks.

Methodology: GLiBRL

The authors propose GLiBRL (Generalised Linear Models in Deep Bayesian RL with Learnable Basis Functions), a framework that enables fully tractable, exact Bayesian inference while maintaining the representational capacity of deep learning.

Core Architecture

GLiBRL enforces a generalized linear relationship between task parameters and data features. Instead of applying neural networks directly to the joint parameter-data space, it separates them:

• Learnable Basis Functions: Neural networks ($\phi_T, \phi_R$) act as expressive, learnable basis functions that map raw state-action data into a feature space ($C_T, C_R$).
• Linear Parameterization: The transition and reward models are linear in these learned features with respect to the task parameters ($\theta_T, \theta_R$).
• Conjugate Priors: The method employs Normal-Wishart priors for the task parameters (mean and noise covariance). This choice is conjugate to the Matrix-Normal likelihood, allowing for exact posterior updates without variational approximation.

Key Technical Components

1. Exact Posterior Updates:

   • The posterior over task parameters is updated sequentially as new samples arrive.
   • By utilizing the Matrix Inversion Lemma, the computational complexity of online updates is reduced to $O(\max(D_S^2, D_T^2))$ for transitions and $O(D_R^2)$ for rewards, ensuring scalability.
   • Crucially, GLiBRL performs exact Bayesian inference over the model noise (covariance matrices $T_\sigma, R_\sigma$), generalizing previous work (like ALPaCA) that assumed known noise.

2. ELBO-Free Objective:

   • Because the posterior is tractable, GLiBRL computes the exact marginal log-likelihood in closed form.
   • The model parameters ($\phi_T, \phi_R$) are optimized by maximizing this marginal likelihood (Equation 19), completely bypassing the need for Variational Inference and the Evidence Lower Bound (ELBO). This eliminates issues like posterior collapse.

3. Permutation Invariance:

   • The exact Bayesian update depends only on the sufficient statistics of the data (the feature matrices and target values), not the order of samples.
   • This permutation invariance allows GLiBRL to seamlessly integrate with both on-policy (e.g., PPO) and off-policy (e.g., SAC) RL algorithms, a significant advantage over methods relying on sequential history encoders.

4. Structural Correspondence:

   • The paper establishes a closed-form identity relating the L2 distance between task representations (derived from posterior means) and the kernel-based empirical correspondence (specifically, signed-measure Cosine Mean Similarity and Maximum Mean Discrepancy) between task samples in the learned feature space.
   • This ensures that task representations are not arbitrary embeddings but are structurally tied to the differences in the underlying data samples.

Key Contributions

1. Tractable Deep BRL: Introduction of a framework that achieves exact Bayesian inference over task parameters and model noise in deep RL settings, eliminating the need for variational approximations.
2. Learnable Basis Functions: A novel architecture that uses neural networks to learn feature bases while preserving the mathematical tractability of linear Bayesian models.
3. Permutation-Invariant Integration: The first deep BRL method capable of seamless integration with off-policy algorithms like SAC due to its inherent permutation invariance.
4. Theoretical Structural Result: The derivation of the first closed-form relationship in online deep BRL/Meta-RL linking the geometric distance of task representations to the kernel-based similarity of task samples.
5. Noise Inference: Explicit Bayesian inference over model noise, which the authors demonstrate reduces predictive errors in both transition and reward models compared to methods assuming fixed noise.

Experimental Results

GLiBRL was evaluated on the MuJoCo locomotion benchmark and the MetaWorld (ML10 and ML45) manipulation benchmark.

• MuJoCo: GLiBRL (combined with SAC) achieved state-of-the-art performance, outperforming PPO-based baselines (RL2, TrMRL, VariBAD) by up to 1.5× in return while using only 10% of their training steps ($10^6$ vs $10^7$). It also outperformed the SAC-based PEARL method by up to 1.8×.
• MetaWorld: In the challenging ML10 and ML45 settings, GLiBRL (combined with PPO) matched or exceeded the performance of strong baselines including MAML, RL2, and VariBAD. Specifically, it achieved up to a 1.1× improvement in success rate on MetaWorld.
• Qualitative Analysis: Visualizations of task representations (via PCA) showed that GLiBRL successfully disentangles continuous task parameters (e.g., target velocities) into a structured latent space after minimal training ($5 \times 10^5$ steps), whereas strong baselines failed to achieve meaningful separation even with significantly larger training budgets.
• Ablation Studies: Removing the inference over model noise (GLiBRL_wo_NI) resulted in higher prediction errors for both transitions and rewards, validating the necessity of the Normal-Wishart prior for accurate modeling.

Significance and Claims

The paper claims that GLiBRL represents a significant advancement in deep Bayesian RL by resolving the tension between the flexibility of deep learning and the rigor of Bayesian inference.

• Theoretical Rigor: It provides the first structural result in online deep BRL linking representation geometry to sample similarity, ensuring that the learned task embeddings are theoretically grounded.
• Practical Efficiency: By avoiding variational inference, the method avoids the instability and optimization difficulties associated with ELBO maximization, leading to more stable learning and better sample efficiency.
• Algorithmic Compatibility: The permutation-invariant nature of the exact update allows for the use of highly efficient off-policy algorithms (SAC), which were previously difficult to integrate with deep BRL methods.
• Generalization: The ability to infer model noise and learn basis functions allows the agent to generalize effectively to unseen tasks with unknown transition and reward dynamics, outperforming current state-of-the-art Meta-RL methods on standard benchmarks.

The authors conclude that while GLiBRL is currently model-free in its policy learning, its learned transition and reward models make it well-suited for future integration with model-based planning, provided computational bottlenecks in sampling from high-dimensional Wishart distributions are addressed.