PAI Course Notes: Probabilistic AI and Uncertainty

The Problem

Given a dataset with input-output pairs, we need to:

• Implement Bayesian Linear Regression to model the relationship between inputs and outputs
• Perform model selection by comparing different polynomial feature transformations
• Use the marginal likelihood (evidence) to select the best model complexity
• Quantify prediction uncertainty using the posterior predictive distribution

Our Approach

We tackled this problem using a fully Bayesian framework:

• Feature Engineering: Applied polynomial transformations of varying degrees (1 to 10) to capture non-linear relationships
• Bayesian Inference: Instead of point estimates, we computed full posterior distributions over model parameters using conjugate priors (Gaussian-Gamma)
• Model Evidence: Calculated the marginal likelihood for each polynomial degree, which naturally penalizes overfitting through the Bayesian Occam's Razor principle
• Posterior Predictive: Generated predictions with credible intervals that reflect both aleatoric (data noise) and epistemic (parameter) uncertainty

Key Insights

• Marginal Likelihood as Model Selection: The evidence automatically trades off model fit and complexity, selecting simpler models when data doesn't justify complexity
• Uncertainty Quantification: Unlike traditional regression, Bayesian methods provide confidence intervals that widen in regions with sparse data
• Occam's Razor in Action: More complex models are penalized unless they significantly improve fit, preventing overfitting
• No Cross-Validation Needed: The marginal likelihood handles model selection in one pass, without requiring train/validation splits

The Problem

Implement Gaussian Process (GP) regression to model complex, non-linear functions with uncertainty quantification:

• Build a Gaussian Process model from scratch using kernel functions
• Handle noisy observations while maintaining smooth predictions
• Implement different kernel functions (RBF, Matérn, Periodic) to capture various function behaviors
• Optimize hyperparameters (length scale, variance, noise) using maximum likelihood estimation
• Provide uncertainty bounds that adapt to data density

Our Approach

We implemented a complete Gaussian Process regression framework:

• Kernel Design: Implemented multiple kernel functions including RBF (Radial Basis Function), Matérn, and Periodic kernels to model different smoothness and periodicity assumptions
• GP Posterior Computation: Derived and implemented the exact posterior distribution over functions given observed data, using the kernel trick to avoid explicit feature space computation
• Hyperparameter Optimization: Maximized the log marginal likelihood with respect to kernel hyperparameters using gradient-based optimization (e.g., L-BFGS)
• Numerical Stability: Applied Cholesky decomposition for efficient and numerically stable inversion of covariance matrices
• Predictive Distribution: Computed mean and variance of the posterior predictive distribution at test points, providing both point predictions and uncertainty estimates

Key Insights

• Non-Parametric Flexibility: GPs define distributions over functions rather than parameters, allowing them to adapt to arbitrary function complexity without fixed model structure
• Kernel as Prior: The choice of kernel encodes our prior beliefs about function smoothness, periodicity, and correlation structure—critical for good predictions
• Uncertainty-Aware Predictions: Predictive variance naturally increases in regions far from training data, providing honest uncertainty quantification
• Computational Challenges: GP inference scales as O(n³) due to matrix inversion, requiring approximations (e.g., inducing points) for large datasets
• Automatic Relevance Determination: Hyperparameter optimization naturally performs feature selection by learning which input dimensions are most relevant
• Connection to Bayesian Linear Regression: GPs can be seen as infinite-dimensional extensions of Bayesian linear regression with infinite basis functions

The Problem

Optimize a black-box objective function while respecting safety constraints—a critical challenge in real-world applications like drug design:

• Maximize bioavailability f(x) of a drug formulation
• Ensure surface area v(x) ≤ 4.0 to remain safe for patients
• Handle noisy observations from expensive experiments
• Never violate safety constraints during the optimization process (safe exploration)
• Balance exploration vs. exploitation within the safe set

The key challenge: we cannot afford to test unsafe configurations, so we must learn both the objective and constraint functions conservatively.

Our Approach

We implemented a dual-GP safe Bayesian optimization framework:

• Dual Gaussian Processes: Maintained two independent GPs—one modeling the objective f(x) and another modeling the constraint violation g(x) = v(x) - 4.0
• Conservative Safe Set: Defined a point as safe if μ_g(x) + β·σ_g(x) ≤ -ε, where β = 3.0 provides high-probability safety guarantees (analogous to UCB but for constraints)
• Kernel Engineering: Used Matérn kernels for f(x) with fixed noise (σ_f = 0.15) and a composite kernel (Linear + Matérn + RBF) for v(x) to capture both global trends and local variations
• Safe Expected Improvement: Maximized Expected Improvement (EI) acquisition function only over the safe set, ensuring all proposed experiments respect safety constraints
• Exploration Bonus: Added a small variance term to the acquisition to avoid premature convergence and encourage safe exploration
• Graceful Degradation: Implemented fallback mechanisms when the GP becomes overly pessimistic—re-sample near previously safe observations

Key Insights

• Safety Through Uncertainty: High uncertainty in constraint predictions forces conservative behavior—the algorithm won't risk unsafe regions until it has more data
• Exploration-Safety Tradeoff: Safe BO must balance three objectives: maximize reward, explore the space, and never violate constraints—a harder problem than unconstrained BO
• Dual Modeling: Separate GPs for objective and constraint allow independent uncertainty quantification, crucial since we need conservative constraint estimates but optimistic objective estimates
• Never Leave the Safe Set: Unlike penalty methods or barrier functions, safe BO provides hard guarantees: if the initial point is safe and the GP is well-calibrated, all queries remain safe with high probability
• Beta Parameter Tuning: The safety parameter β controls risk aversion—higher β means more conservative (larger safe set margin), lower β allows more aggressive exploration
• Real-World Applications: This framework is critical for domains where constraint violations are unacceptable: medical treatments, autonomous systems, chemical processes, robotics
• Sample Efficiency: Safe BO achieves near-optimal solutions in fewer evaluations than random search or grid search while maintaining safety, crucial when experiments are expensive or time-consuming

The Problem

Implement a state-of-the-art deep reinforcement learning algorithm for continuous control:

• Train an agent to solve the CartPole-v1 control task using only raw observations
• Handle continuous action spaces with Gaussian stochastic policies
• Balance exploration and exploitation while learning from experience
• Ensure stable learning by constraining policy updates to avoid catastrophic forgetting
• Use off-policy learning with a replay buffer for sample efficiency
• Implement actor-critic architecture with twin Q-functions for value estimation

Our Approach

We implemented MPO, an EM-style algorithm that constrains policy updates using KL divergence:

• Twin Critics (Clipped Double Q-Learning): Maintained two independent Q-networks (Q₁, Q₂) and used their minimum for target computation to reduce overestimation bias
• Gaussian Stochastic Actor: Policy network outputs mean μ and log-std for each action dimension, creating a diagonal Gaussian distribution. Actions are sampled using the reparameterization trick and squashed through tanh for bounded control
• E-Step (Critic Update): Updated critics toward Bellman targets computed by sampling K actions from the target policy at next states and averaging their Q-values
• M-Step (Policy Update): Updated policy to maximize expected Q-value using importance-weighted maximum likelihood:
  • Sample K actions from current policy at observed states
  • Compute importance weights w_k ∝ exp(Q(s,a_k)/η) using softmax
  • Maximize weighted log-likelihood: Σ_k w_k log π(a_k|s)
• Temperature Parameter η: Adaptively controlled via Lagrangian optimization to maintain KL(π_old || π_new) ≈ ε_KL, preventing drastic policy changes
• Soft Target Updates: Applied Polyak averaging (τ=0.005) to slowly update target networks, stabilizing learning
• Replay Buffer: Stored 50k transitions for off-policy sampling, enabling data reuse and breaking temporal correlations

Key Insights

• EM Perspective on RL: MPO frames policy improvement as an EM algorithm—E-step evaluates actions, M-step improves policy toward high-value actions weighted by their advantage
• Trust Region via KL Constraint: Unlike PPO's clipping or TRPO's line search, MPO uses a temperature parameter to implicitly enforce trust regions, ensuring smooth policy evolution
• Sample Reweighting: Importance weights naturally implement advantage-weighted regression: actions with high Q-values get more influence on policy updates
• Continuous Control Challenges: Tanh squashing ensures actions stay within bounds while maintaining differentiability for gradient-based optimization
• Twin Q-Functions: Taking min(Q₁, Q₂) for target computation reduces overoptimistic Q-value estimates that plague single-critic methods, improving stability
• Off-Policy Efficiency: Replay buffer enables multiple gradient updates per environment step, dramatically improving sample efficiency compared to on-policy methods like PPO
• Reparameterization Trick: Sampling a ~ N(μ, σ) as a = μ + σ·ε (ε ~ N(0,1)) allows backpropagation through stochastic sampling, critical for policy gradient estimation
• Stabilization Techniques: Combination of target networks, twin critics, soft updates, and KL constraints creates a remarkably stable learning algorithm suitable for complex control tasks