On the Role of Computation in Reinforcement Learning

Technical Summary: On the Role of Computation in Reinforcement Learning

Problem Statement
The standard view in Reinforcement Learning (RL) treats policies as static functions that map states to actions, spending a fixed amount of computation regardless of the complexity of the decision required. This paradigm conflates the number of parameters (memory) with the amount of computation (time). While scaling parameter counts (Moore's Law) has been a common strategy, the compute required to make optimal decisions can vary by orders of magnitude across different states. Static, massive policies are often inefficient, and computationally limited policies may fail to generalize to longer-horizon tasks, instead learning heuristics that overfit to training distributions. The paper addresses the lack of a formal language to distinguish between compute and parameters in RL and investigates whether policies with fixed parameters can benefit from variable, increased computation time.

Methodology
The paper approaches the problem through two complementary lenses: theoretical formalization and empirical validation.

1. Theoretical Framework:

   • The authors formalize RL policies as bounded models of computation using Turing machines.
   • They define time-bounded policies ($\Pi_t$) as those implementable by a Turing machine halting within $t(|s|)$ steps.
   • They introduce the Policy Hierarchy Theorem, proving that for time-constructible functions $g(n) \in o(t(n)/\log t(n))$, there exist Markov Decision Processes (MDPs) where policies with higher compute budgets ($\Pi_t$) perform arbitrarily better than those with lower budgets ($\Pi_g$).
   • They further prove the Long Horizon Generalization Theorem, demonstrating that low-compute policies can overfit to training tasks (matching the optimal value function on training data) but fail to generalize to longer-horizon test tasks, whereas high-compute policies can solve these extended tasks.

2. Architectural Proposal:

   • To test these theories, the authors propose a minimal recurrent architecture called the Interpolation Recurrent Unit (IRU).
   • The IRU is a simplified variant of LSTM/GRU blocks, consisting of two linear layers (Forget and Input) that interpolate between the previous cell state and a new local computation.
   • The architecture applies this block $N$ times during a forward pass (recurrent steps), allowing the policy to use a variable amount of compute ($N$) while maintaining a fixed parameter count.
   • The full network (IRU-$N$) includes an initial linear layer, $N$ recurrent blocks with skip connections, and a final output head.

3. Experimental Setup:

   • Experiments were conducted across 31 tasks spanning discrete (Boxpick, Lightsout) and continuous (OGBench) domains, including both online and offline RL settings.
   • Baselines included standard Feedforward Networks (MLP) and Residual Networks (ResNet). Crucially, the ResNet baseline was designed to have a similar effective depth (compute) but significantly more parameters (up to 5x), while the MLP had similar parameters but less compute.
   • The study evaluated performance on training tasks and, critically, on unseen longer-horizon test tasks to measure generalization.

Key Contributions

• Formalization of Compute-Bounded Policies: The paper provides the first formal proof that increasing computation time (distinct from parameter count) allows policies to solve problems and generalize to longer horizons that are theoretically impossible for lower-compute policies.
• Minimal Recurrent Architecture: The introduction of the IRU, a lightweight architecture capable of scaling compute via recurrent steps without increasing parameter count.
• Value of Compute Metric: A proposed numerical measure, the Value of Compute (VoC), analogous to the Value of Information in POMDPs. This metric estimates the gain or loss in value when an agent temporarily increases its compute budget for a specific number of steps.
• Disentangling Compute and Parameters: Empirical evidence showing that performance gains come from increased expressivity (compute) rather than just increased model capacity (parameters).

Results

• Performance Scaling: In both discrete and continuous environments, increasing the number of recurrent steps ($N$) from 1 to 10 significantly improved success rates. In challenging tasks like Boxpick and OGBench, performance increased up to 10-fold.
• Generalization: IRU policies trained on short-horizon tasks generalized significantly better to unseen long-horizon tasks compared to MLPs and ResNets. Notably, IRU outperformed ResNets with 5 times more parameters on long-horizon evaluation tasks.
• Sample Efficiency: IRU architectures demonstrated higher sample efficiency, converging to higher asymptotic performance with fewer environment steps.
• Source of Improvement: "Yoked" experiments (training a passive agent on data collected by an active agent) revealed that the primary benefit of additional compute stems from the expressivity of the policy (its ability to learn complex solutions from data) rather than improved exploration or data quality. The passive agent trained on data from a low-compute active agent performed poorly unless it itself used high compute.
• Ablation: The IRU architecture outperformed standard LSTMs and Recurrent ResNets in terms of speed and parameter efficiency, suggesting that specific gating mechanisms are necessary to leverage the benefits of additional compute effectively.

Significance and Claims
The paper claims to shift the perspective of RL from static function approximation to a computational view where time and memory are distinct axes. The authors argue that:

1. Compute is a distinct resource: Increasing computation time can provably improve performance and generalization without increasing the hypothesis class size (parameters), thereby avoiding the overfitting issues associated with larger static models.
2. Generalization via Iteration: Iterative computation allows policies to learn algorithmic solutions that extrapolate to longer horizons, a phenomenon observed in language models but previously underexplored in RL without pre-training or multi-task curricula.
3. Limitations of Static Views: The results suggest that standard MDP assumptions (e.g., the existence of a deterministic optimal policy) may not hold under computational constraints, potentially explaining why non-Markov policies are sometimes necessary even in fully observed tasks.

The authors conclude that while their work uses a minimal architecture to isolate these effects, the findings suggest that future RL systems should consider adaptive compute strategies, where the amount of computation is dynamically scaled based on the difficulty of the current state.