PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexLoop: Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL

Xun Wang, Ruishuo Chen, Yu Chen, Zhuoran Li, Longbo Huang

arXiv:2609.34488v1cs.LGcs.AI

TL;DR

Pretrained looped RL policies trained at fixed depth concentrate reliable decisions near that depth, leaving reliable inference across recurrent depths an open challenge. FlexLoop post-trains pretrained policies with the original RL objective and adjacent-depth policy distillation to transfer decision quality from deeper to shallower steps. Across 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth performance and enables reliable shallower-depth inference, reducing average recurrent computation by up to 43%.

  • Problem

    Pretrained looped RL policies trained at fixed depth concentrate reliable decisions near that depth, leaving reliable inference across recurrent depths an open challenge.

    Source: §1 · p. 2

  • Method

    FlexLoop post-trains pretrained policies with the original RL objective and adjacent-depth policy distillation to transfer decision quality from deeper to shallower steps.

    Source: §1 · p. 2

  • Results

    Across 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth performance and enables reliable shallower-depth inference, reducing average recurrent computation by up to 43%.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Consistency across successive recurrent depths provides a signal for state-wise adaptive inference that can stop computation once the policy stabilizes.

    Source: §4.1 · p. 5

    In offline settings, the agent learns from a fixed dataset without further environment interaction.

    Source: §3 · p. 3

Abstract

from arXiv · show

Looped architectures scale computation by reusing the same parameters across recurrent steps, and recent work shows that they substantially improve deep reinforcement learning policies on long-horizon tasks. Since recurrent depth directly controls computation, one may expect looped policies to naturally support elastic inference across recurrent depths. Surprisingly, we find that pretrained looped policies exhibit severe recurrent-depth specialization: reliable decisions are concentrated near the full trained depth, tying deployment computation to this depth even when less computation may suffice. Achieving depth elasticity, i.e., reliable decisions across recurrent depths with adaptive computation at deployment, therefore remains a key challenge. To address this, we propose FlexLoop, a novel post-training framework that converts pretrained fixed-depth looped policies into depth-elastic policies. FlexLoop keeps training on the original RL objective to preserve full-depth capability while performing adjacent-depth policy distillation to progressively transfer decision quality from deeper to shallower recurrent steps. The resulting policy supports reliable inference across recurrent depths and enables state-wise adaptive inference through recurrent-depth consistency. Experiments on $30$ online and offline long-horizon goal-conditioned environments show that FlexLoop preserves full-depth performance while making shallower depths effective. Keeping competitive performance, FlexLoop reduces average recurrent depth by up to $\bf{43\%}$ and achieves up to $\bf{1.34\times}$ wall-clock speedup in a stress test.

1. Introduction

Looped policies offer a way to scale recurrent computation, but fixed-depth training concentrates reliable decisions near the trained depth. FlexLoop combines continued RL optimization with adjacent-depth distillation; experiments report reliable shallower-depth inference and reduced deployment computation.

1. Introduction

  • Looped architectures reuse computation to increase effective depth without increasing trainable parameters, making recurrent depth a potential deployment-time computation budget.Prior work found that additional recurrent steps can improve long-horizon planning and decision making.Source: §1 · pp. 1–2
  • Jointly training multiple depths can suppress full-depth performance, while distillation alone can degrade generalization to unseen tasks.These empirical challenges motivate FlexLoop’s combination of continued RL optimization and distillation.Source: §1 · p. 2
  • Consistency across recurrent depths provides a signal for state-wise adaptive inference at deployment.Computation can stop once the policy has sufficiently stabilized.Source: §1 · pp. 2–3
  • Fixed-depth training concentrates reliable decisions near the full trained depth, limiting adaptation to deployment computation constraints.The paper asks whether full-depth capability can be preserved while making shallower recurrent depths reliable.Source: §1 · p. 2
  • FlexLoop post-trains a pretrained looped policy with the original RL objective and adjacent-depth distillation, transferring decision quality from deeper to shallower steps.The objective preserves full-depth capability while the distillation progressively propagates decision quality toward shallower recurrent depths.Source: §1 · p. 2
  • Across 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth performance and makes shallower depths reliable, reducing average recurrent depth by up to 43%.It achieves up to 1.34× wall-clock speedup in a stress test while maintaining competitive or better task performance.Source: §1 · pp. 2–3

2. Related Work

Prior work explores added computation in deep RL and variable inference depth in looped models. The paper investigates depth elasticity in looped RL, where reliable intermediate decisions remain underexplored.

Computation Scaling in DRL

  • Recurrent computation has been shown to scale RL capability without increasing policy parameter count, and this work examines whether policies remain effective across inference depths.The authors report that standard fixed-depth training instead leads to recurrent-depth specialization.Source: §2 · p. 3

Looped Architectures and Elastic Depth

  • Looped models reuse parameters across computation steps, and prior work studies variable inference budgets, self-distillation, and adaptive stopping in language and visual-generation models.Examples include LoopFormer, ELT, and adaptive inference methods for looped language models.Source: §2 · p. 3
  • Reliable depth elasticity in RL remains underexplored because RL lacks ground-truth intermediate decisions and policy quality is determined by long-horizon return rather than local output agreement.The passage contrasts RL with settings that have explicit targets.Source: §2 · p. 3

3. Preliminary

The preliminary section defines the goal-conditioned RL setting and describes looped policies and IRU. It frames the objective as making intermediate-depth outputs reliable while retaining full-depth capability.

Goal-Conditioned RL

  • A goal-conditioned Markov decision process specifies state, action, and goal spaces, transition dynamics, reward, an initial state-goal distribution, and a discount factor.The policy conditions on state and goal to generate actions.Source: §3 · p. 3

Looped Policies and IRU

  • Looped architectures increase computational depth by repeatedly applying a shared computation block without introducing depth-specific parameters.The representation after k recurrent steps is denoted h_k, and the shared block is F_θ.Source: §3 · pp. 3–4
  • The paper focuses primarily on IRU-based policies, whose depth-k output may represent action values, a policy, or a value estimate.Standard IRU policies are optimized through full depth K; FlexLoop aims to make outputs at k < K reliable while preserving full-depth capability.Source: §3 · p. 4

4. Method: FlexLoop

FlexLoop post-trains a pretrained looped policy with the original RL objective and adjacent-depth policy distillation, preserving full-depth capability while making shallower decisions reliable. It instantiates the approach across three goal-conditioned RL algorithms and uses adjacent-depth policy discrepancy to guide adaptive inference.

4.1. Elastic Post-Training

  • FlexLoop combines the original full-depth RL objective with progressive adjacent-depth distillation to preserve pretrained capability while transferring decision quality to shallower policies.The discrepancy between successive policies can also signal when a policy has stabilized enough to stop computation.Source: §4.1 · pp. 4–5

4.2. Instantiating FlexLoop across RL Algorithms

  • FlexLoop is instantiated for goal-conditioned DQN, PPO, and IQL across online and offline settings and discrete and continuous action spaces.GC-IQL distills Gaussian policies using normalized policy KL to avoid the scale-inflation degeneracy of raw KL.Source: §4.2 · pp. 5–6

4.3. State-Wise Adaptive Inference

  • The largest adjacent-depth policy discrepancy bounds the performance gap between any shallow policy and the full-depth policy under the proposition’s stated assumptions.A separate proposition establishes that strong full-depth performance alone does not guarantee useful intermediate decisions.Source: §4.3 · p. 6
  • Smaller local discrepancies are associated with smaller divergence from the full-depth policy on training and unseen tasks, motivating discrepancy-threshold stopping.Changing the threshold yields different computation–performance trade-offs from one trained policy.Source: §4.3 · p. 7, Figure 3 · p. 7

5. Experiments

Experiments test FlexLoop on 30 long-horizon goal-conditioned environments spanning online and offline RL, multiple action spaces, and unseen-task generalization. FlexLoop improves shallow-depth decisions while preserving full-depth capability, supports adaptive computation, and is examined across training choices and architectures.

5.1. Experimental Setup

  • Evaluation covers 30 long-horizon goal-conditioned environments across online and offline RL, discrete and continuous actions, and in-distribution and out-of-distribution settings.The experiments use pretrained IRU-5 policies and retain the baseline environment configurations, RL algorithms, and evaluation protocols unless otherwise stated.Source: §5.1 · p. 7

5.2. Analysis of Learned Depth Elasticity; 5.3. State-wise Adaptive Inference

  • FlexLoop improves shallow-depth decisions while maintaining or improving full-depth performance on training tasks and unseen tasks.For adaptive inference, it reduces average recurrent depth by up to 43% while maintaining full-depth policy performance, and achieves up to 1.34× wall-clock speedup compared with fixed-depth inference.Source: §5.2 · p. 8, Table 1 · p. 9, §5.3 · pp. 9–10

5.2. Analysis of Learned Depth Elasticity

  • Shallower recurrent depths improve more slowly, and FlexLoop’s gains are largest on GC-DQN BoxPick and smaller on PPO and offline IQL tasks.The authors hypothesize that value-based action selection makes depth-wise adaptation more direct than separately parameterized policies.Source: §5.2 · pp. 8–9

5.3. State-wise Adaptive Inference

  • Varying the stopping threshold trades computation for performance, and the favorable trade-off also holds on unseen tasks.Smaller thresholds generally allocate deeper computation and achieve higher performance; larger thresholds reduce average inference depth with controllable trade-offs.Source: §5.3 · p. 9

5.4. Ablation Study

  • Ablations show that continued RL optimization preserves full-depth capability and generalization, while depth-wise distillation transfers decision quality to shallower depths.Self-TD alone yields limited shallow-depth improvement; AnyK-TD and Distill Only each have limitations in full-depth capability or unseen-task generalization.Source: §5.4 · p. 10
  • Increasing λdepth improves shallow-depth performance on training tasks, but excessive distillation can hurt unseen-task generalization.A moderate λdepth balances shallow-depth adaptation and generalization.Source: §5.4 · p. 10
  • FlexLoop improves shallow-depth performance and preserves full-depth capability in both LSTM variants, with greater gains for LSTM-256 than LSTM-184.LSTM-184 has a parameter count comparable to IRU, while LSTM-256 has larger capacity.Source: §5.4 · p. 10

5.5. Visualization of Adaptive Inference

  • On successful OGBench trajectories, FlexLoop allocates deeper computation near difficult decisions and terminates earlier after adjacent-depth policies stabilize.Examples include deeper computation at skill boundaries in Scene and near turns and goals in AntMaze, versus two recurrent steps on PointMaze.Source: Figure 6 · p. 11, §5.5 · p. 11

6. Conclusion and Future Directions

The paper identifies recurrent-depth specialization and presents FlexLoop as post-training with continued RL optimization and adjacent-depth distillation. Across 30 environments, it reports reliable shallower-depth inference and lower computation, while outlining several directions for future work.

6. Conclusion and Future Directions

  • Across 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth capability and enables reliable shallower-depth inference, with up to 43% less average recurrent computation and 1.34× wall-clock speedup.Task performance remains competitive or better.Source: §6 · p. 11
  • FlexLoop induces depth elasticity through continued RL optimization and adjacent-depth policy distillation.The paper identifies recurrent-depth specialization in pretrained looped RL policies.Source: §6 · p. 11
  • Future work includes extending depth elasticity beyond the trained depth and developing adaptive-inference mechanisms beyond the current depth-consistency criterion.Source: §6 · p. 11
  • The authors also propose studying depth elasticity in larger-scale looped architectures, including transformer-based models.Source: §6 · p. 11

A.1. Proof of Proposition 4.1

The appendix proves that shared-weight recurrent policies can behave very differently at full and shallower depths, outlines a return-analysis proof, and describes the evaluated environment suites.

A.1. Proof of Proposition 4.1

  • A constructed shared-weight recurrent policy can be arbitrarily close to optimal at full depth while its shallower policies are arbitrarily close to the minimum achievable return.The construction uses a recurrent state with h_k=k while the recurrent block ignores the state and goal.Source: §A.1 · p. 15

A.2. Proof of Proposition 4.2

  • The second proof relates adjacent-depth total-variation policy discrepancies to return using the performance difference lemma and the triangle inequality across depths.Its analysis uses an equivalent discounted infinite-horizon formulation, with time added to the state when needed and x=(s, g).Source: §A.2 · pp. 15–16

B.1. Environments

  • The evaluation covers 30 long-horizon goal-conditioned environments: discrete-action BoxPick and LightsOut tasks, plus continuous-action offline OGBench scenarios.BoxPick tests box transport and generalization; LightsOut includes unseen configurations requiring longer solutions; OGBench covers navigation and manipulation scenarios.Source: §B.1 · pp. 16–17

B.2. Implementation Details

Implementation follows the IRU codebase and generally adopts its training settings, with specified environment-specific choices. The extended-results section supplies environment-wise evaluations alongside the paper's aggregate results.

B.2. Implementation Details

  • The implementation reproduces the IRU baseline and generally uses its training steps, environment configurations, and hyperparameters, with stated exceptions for OGBench and on-policy GC-PPO.OGBench uses hidden size 256 rather than 64; GC-PPO replaces GAE-based value targets with pretrained IRU value estimates to stabilize FlexLoop training, while retaining standard PPO policy updates.Source: §B.2 · p. 17
  • The adjacent-depth policy distillation weight λ_depth is FlexLoop’s key hyperparameter, with environment-specific values reported in Table 3.Source: §B.2 · p. 17

C. Extended Results

  • The extended results provide environment-wise evaluations that complement the aggregate results in the main paper.Source: C. Extended Results · p. 18

C.1. Full training curves for Figure 4

The training-curve figures report improved shallow-depth performance while retaining full-depth capability, and the elasticity curves compare pretrained and post-trained policies on training and unseen tasks.

C.1. Full training curves for Figure 4

  • On unseen tasks, FlexLoop improves shallow-depth decision elasticity while maintaining out-of-distribution generalization.Figure 7 compares pretrained IRU-5 and FlexLoop depth-wise mean success rates with one-standard-deviation shading.Source: Figure 7 · p. 18
  • During pretraining and post-training, FlexLoop rapidly improves shallow-depth performance while preserving pretrained full-depth capability.Figure 8 plots depth-wise mean success rates with one-standard-deviation shading.Source: Figure 8 · p. 19

C.2. Elasticity Curves

  • Figures 9 and 10 present depth-wise decision-performance curves before and after FlexLoop post-training, with Figure 9 showing training-task results and Figure 10 corresponding unseen-task results.Markers denote mean success rates, and shading indicates one standard deviation.Source: §C.2 · p. 20, Figure 9 · p. 20, Figure 10 · p. 21

C.3. Full results for Table 1

The section presents full state-wise adaptive inference results, explains how Pareto curves are constructed, and illustrates how FlexLoop varies recurrent depth during a Scene trajectory.

C.3. Full results for Table 1

  • Table 4 reports full state-wise adaptive inference results across all environments, with ± indicating one standard deviation.Source: Table 4 · p. 22

C.4. Pareto Curve

  • Pareto curves plot task performance against average recurrent depth, with each point obtained by sweeping the stopping threshold ε.The ε grids vary by environment family; manipulation tasks use finer grids than navigation tasks in OGBench.Source: Figure 11 · p. 23, §C.4 · p. 23

C.5. Visualized Trajectory Segments of Scene

  • In Scene trajectories, FlexLoop uses fewer recurrent steps for simple, monotonic motions and more computation at challenging decisions such as goal switching.The figure shows selected segments from a single episode.Source: §C.5 · p. 24, Figure 12 · p. 24

D. Additional Ablations

The additional ablations examine learning FlexLoop from queried teacher outputs and applying post-training to policies with different pretrained recurrent depths. They report elasticity and full-depth capability across these settings.

D.1. FlexLoop from Scratch

  • When pretrained weights are unavailable, the BoxPick-Exact-4 GC-DQN setting tests whether random initialization can acquire elasticity and recover the teacher’s full-depth capability.Only outputs of the pretrained policy can be queried.Source: §D.1 · p. 24
  • The scratch modifications replace the full-depth standard TD target with a queried teacher value and add an inter-depth Bellman refinement objective for shallow steps.Source: §D.1 · p. 24
  • These modifications compensate for missing pretrained initialization and provide no additional benefit in the teacher-initialized setting.Source: §D.1 · p. 25
  • From scratch, FlexLoop rapidly develops depth elasticity, recovers strong full-depth performance, and supports state-wise adaptive inference.Figures 13 and 14 show results on training and unseen tasks; this setting uses output-level access to the pretrained policy.Source: §D.1 · p. 25, Figure 13 · p. 25, Figure 14 · p. 25

D.2. FlexLoop on Policies with Different Pretraining Depths

  • Across all pretraining depths, FlexLoop substantially improves shallow-depth performance while preserving full-depth capability, and the same trend holds on unseen tasks; state-wise adaptive inference remains effective across depths.Deeper pretrained policies generally require smaller stopping thresholds ε.Source: §D.2 · p. 25

Related papers

Loading 2609.34488v1…