Scaling Reinforcement Learning with Composite Gradients

Peter Bubenik · Sony AI · · Source

After studying this material, students should be able to:

  1. Explain the two main approaches to policy gradient estimation in MBRL
  2. Understand why combining estimators (composite methods) can outperform individual methods
  3. Describe how Inverse Variance Weighting (IVW) works conceptually
  4. Understand what TPX improves over TP and why it matters
  5. Recognize the practical significance of scalability in modern RL tasks

Step-by-Step Teaching

Step 1: Foundation — What is Model-Based Reinforcement Learning (MBRL)?

What is Reinforcement Learning?

Agent → takes Action → Environment → gives Reward + New State
         ↑___________________________________|

An agent learns a policy (strategy) to maximize cumulative rewards.

What makes it "Model-Based"?

Instead of only learning from real experience:

ApproachHow it learns
Model-Free RLOnly real environment interactions
Model-Based RLLearns a model of the environment, then simulates/plans with it

Key Insight: A learned model lets the agent "imagine" outcomes without costly real-world interactions.

Why does this matter?

  • Real interactions can be expensive (robotics, visual tasks)
  • A model allows planning ahead through simulated trajectories

Step 2: The Core Problem — How Do We Compute Policy Gradients?

To improve a policy, we need to compute gradients — the direction to adjust policy parameters to get more reward.

The Challenge:

Policy Parameters (θ) → Actions → Rewards

We need: "How does changing θ affect total reward?"

There are two fundamentally different approaches:


Step 3: The Two Estimator Types

Approach 1: Likelihood Ratio (LR) Gradients

Also called REINFORCE or "derivative-free"

Core Idea: Don't differentiate through the environment. Instead, observe outcomes and weight them by probability.

Intuition:
- Try many actions
- Actions that led to HIGH reward → increase their probability
- Actions that led to LOW reward → decrease their probability

Formula concept:

∇θ J ≈ E[∇θ log π(a|s) · R]
         ↑                  ↑
    "How likely was    "How good was
     this action?"      the outcome?"

Pros and Cons:

✅ Pros❌ Cons
Works even if model is non-differentiableVery HIGH variance
Simple to implementNeeds many samples to be accurate
General purposeSlow convergence

Approach 2: Reparameterization (RP) Gradients

Also called "pathwise gradients" or backpropagation through time

Core Idea: If the model is differentiable, backpropagate gradients directly through the simulation.

θ → action → model → next state → model → ... → Reward
 ←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←
              Backpropagate gradients directly

Intuition: Like training a neural network — chain rule all the way back.

Pros and Cons:

✅ Pros❌ Cons
LOW varianceRequires differentiable model
Efficient use of dataGradients can explode or vanish over long horizons
Fast convergenceNumerically unstable for long sequences

Step 4: The Key Insight — Why Not Use Both?

The Variance-Bias Tradeoff Visualization:

LR Estimator:   [============================] High Variance, works anywhere
RP Estimator:   [====]                         Low Variance, but unstable long-term

Combined:       [========]                     Best of both worlds?

Introducing: Composite Estimators

Core Idea: At each step in a trajectory, choose or blend between LR and RP based on which is more reliable.

Trajectory:  s₀ → s₁ → s₂ → s₃ → s₄ → ... → sₙ
              
Use RP:      [←←←←←←←←←]  (recent steps, stable)
Use LR:      [←←←←←←←←←←←←←←←←←←←←←]  (earlier steps, RP unstable)

Key Insight: RP gradients are reliable for short horizons but degrade over long ones. LR can cover the rest.


Step 5: Inverse Variance Weighting (IVW) — The Combining Mechanism

What is IVW?

IVW is a statistical technique to optimally combine multiple estimators of the same quantity.

Intuition: Trust estimators that are more consistent (low variance) more heavily.

Formula:

If estimator 1 has variance σ₁²
If estimator 2 has variance σ₂²

Combined weight for estimator 1: w₁ = (1/σ₁²) / (1/σ₁² + 1/σ₂²)
Combined weight for estimator 2: w₂ = (1/σ₂²) / (1/σ₁² + 1/σ₂²)

Final estimate = w₁·estimate₁ + w₂·estimate₂

Simple Example:

Estimator A: variance = 1    → weight = 1/1 = 1.0  → 50% weight
Estimator B: variance = 0.1  → weight = 1/0.1 = 10 → ~91% weight

→ Trust B much more because it's more consistent!

Why is this optimal?

  • Mathematically proven to minimize variance of the combined estimate
  • Automatically adapts — no manual tuning of blend ratio needed

Step 6: Total Propagation (TP) — Prior Work

What TP Does:

TP applies IVW at every node along the trajectory to blend LR and RP gradients.

Trajectory nodes:
s₀ ——→ s₁ ——→ s₂ ——→ s₃ ——→ Reward
 ↑      ↑      ↑      ↑
IVW    IVW    IVW    IVW
(blend LR+RP at each point)

TP's Achievement:

Orders of magnitude improvement over using either LR or RP alone

TP's Problem:

  • Not scalable to modern, complex RL tasks
  • Difficult to implement efficiently
  • Had not been tested on visual/high-dimensional tasks

Step 7: Total Propagation X (TPX) — The New Contribution

What TPX Improves:

Improvement 1: Different Node for IVW

TP:   Apply IVW at every intermediate node
TPX:  Apply IVW at a strategically chosen node
      → Less computation, same benefit

Improvement 2: Coordinate-wise Weighting

TP:   Single weight for entire gradient vector
      [w · (g₁, g₂, g₃, g₄, g₅)]

TPX:  Different weight per dimension
      [(w₁·g₁, w₂·g₂, w₃·g₃, w₄·g₄, w₅·g₅)]

Why coordinate-wise matters:

  • Policy parameters have many dimensions
  • Some dimensions may have high-variance LR estimates
  • Others may have stable RP estimates
  • Treating them individually → better overall estimate

TPX Summary:

FeatureTPTPX
IVW applicationAll nodesStrategic node
WeightingScalarPer-coordinate
ScalabilityLimited✅ Modern tasks
ImplementationComplex✅ Easier

Step 8: Practical Validation — Dreamer + TPX

What is Dreamer?

Dreamer is a state-of-the-art visual MBRL algorithm:

  • Learns from pixel observations (images)
  • Builds a world model in latent space
  • Plans using imagined trajectories
Real pixels → Encoder → Latent State → World Model → Imagined Future
                                                           ↓
                                                    Policy Gradient

The Experiment Results:

Problem with standard Dreamer:

Short horizon:  ✅ Works well
Long horizon:   ❌ Fails (RP gradients explode/vanish)

With TPX:

Short horizon:  ✅ Works well  
Long horizon:   ✅ Works reliably
Cost:           Only a fraction of additional computation

Step 9: Putting It All Together

Complete Conceptual Map:

MBRL Problem
     │
     ▼
Need Policy Gradients
     │
     ├──→ LR Estimator (high variance, general)
     │
     ├──→ RP Estimator (low variance, unstable long-term)
     │
     └──→ COMBINE THEM (best of both)
               │
               ▼
         IVW Weighting
         (trust low-variance more)
               │
               ▼
         TP (prior work, not scalable)
               │
               ▼
         TPX (scalable, coordinate-wise)
               │
               ▼
         Applied to Dreamer
         ✅ Solves long-horizon failure
         ✅ Minimal extra cost
         ✅ Easy to implement

Step 10: Key Takeaways

ConceptRemember This
LR gradientsHigh variance, works anywhere, derivative-free
RP gradientsLow variance, needs differentiable model, unstable long-term
IVWCombine estimators by trusting lower-variance ones more
TPFirst composite method, proved concept, not scalable
TPXScalable TP with coordinate-wise weighting
Key resultLong horizons that break Dreamer work with TPX

Self-Check Questions

  1. Why does RP gradient estimation fail for long simulation horizons?
  2. What does "inverse" mean in Inverse Variance Weighting — why invert the variance?
  3. Why is coordinate-wise weighting better than scalar weighting?
  4. What specific problem does TPX solve that TP could not?
  5. Why would visual MBRL (like Dreamer) be a good test of scalability?

More to study