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 43% and achieves up to 1.34× wall-clock speedup in a stress test.
Figures & tables
Figure 1 : Fixed-depth training induces recurrent-depth specialization. Each IRU- K is trained at the K -th recurrent depth. Reliable decisions are concentrated near the full trained depth.
Figure 2 : Framework of FlexLoop. (a) Starting from a depth-specialized looped policy, FlexLoop preserves full-depth capability with the original RL objective while propagating decision quality to shallower depths via adjacent-depth policy distillation. (b) The resulting policy supports reliable decisions across recurrent depths. (c) Adjacent-depth policy discrepancy serves as a stopping signal for state-wise adaptive inference, allocating more computation to harder states.
Figure 3 : (a) FlexLoop reduces adjacent-depth policy discrepancy during post-training and across depths. (b) Smaller δk corresponds to smaller DKL(πk∥πK) on both training and unseen tasks. Marker size denotes the sample fraction in each pooled quantile bin.
Figure 4 : Depth-wise decision performance of pretrained IRU- 5 and FlexLoop on training and unseen tasks. Solid lines denote mean success rates and shading indicates one standard deviation.
Environment
FlexLoop
IRU-5
Avg. Depth
Succ.
Succ.
OGBench
4.19 ± 0.08
0.410 ± 0.020
0.385 ± 0.029
BoxPick (Train)
4.50 ± 0.02
0.989 ± 0.002
0.989 ± 0.003
BoxPick (Unseen)
4.46 ± 0.02
0.778 ± 0.055
0.760 ± 0.043
LightsOut-4x5 (Train)
2.84 ± 0.08
0.747 ± 0.081
0.732 ± 0.076
LightsOut-4x5 (Unseen)
2.93 ± 0.21
0.053 ± 0.031
0.013 ± 0.004
Table 1 : State-wise adaptive inference results across environments. ± indicates one standard deviation.
Model
Success
Avg. Depth
Wall-clock (ms)
Train
IRU-5
0.995±0.001
5
170.1±17.2
FlexLoop
0.993±0.001
3.85±0.04
133.0±7.7
IRU-10
0.997±0.000
10
333.2±36.5
FlexLoop
0.996±0.001
9.35±0.04
319.9±22.2
Unseen
IRU-5
0.511±0.085
5
170.1±17.2
FlexLoop
0.543±0.093
3.80±0.05
127.4±6.3
Table 2 : Wall-clock comparison on BoxPick-Exact-4. ± indicates one standard deviation. Bold numbers indicate the shortest wall-clock latency.
Figure 5 : Ablation studies of FlexLoop. Solid lines denote mean performance, while shaded regions indicate one standard deviation.
Figure 6 : State-wise adaptive recurrent-depth allocation across OGBench tasks. FlexLoop allocates more recurrent computation to states requiring more complex decisions.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Environment
λdepth
BoxPick-Exact-3
1.0
BoxPick-Exact-4
1.0
BoxPick-Generalize-4-1
1.0
BoxPick-Generalize-4-2
1.0
LightsOut
0.05
Scene
0.05
Appendix
Table 3 : Adjacent-depth policy distillation weight used for each environment.
Figure 7 : Depth-wise decision performance of pretrained IRU- 5 and FlexLoop on unseen tasks. Solid lines denote mean success rates and shading indicates one standard deviation. FlexLoop improves shallow-depth decision elasticity while maintaining out-of-distribution generalization.
Figure 8 : Depth-wise decision performance during IRU pretraining and FlexLoop post-training. Solid lines denote mean success rates and shading indicates one standard deviation. FlexLoop rapidly improves shallow-depth performance while preserving the full-depth capability of pretrained policies.
Figure 9 : Depth-wise decision performance on training tasks before and after FlexLoop post-training. Markers denote mean success rates and shading indicates one standard deviation.
Figure 10 : Depth-wise decision performance on unseen tasks before and after FlexLoop post-training. Markers denote mean success rates and shading indicates one standard deviation.
Task
FlexLoop
IRU-5
Avg. Depth
Succ.
Succ.
Scene
4.66 ± 0.06
0.496 ± 0.064
0.489 ± 0.087
Cube
4.75 ± 0.03
0.569 ± 0.056
0.492 ± 0.114
Puzzle
3.45 ± 0.03
0.296 ± 0.029
0.293 ± 0.01
AntMaze
4.80 ± 0.03
0.292 ± 0.059
0.253 ± 0.034
PointMaze
3.28 ± 0.30
0.395 ± 0.004
0.396 ± 0.003
Appendix
Table 4 : Full state-wise adaptive inference results on all environments. ± indicates one standard deviation.
Figure 11 : Pareto trade-offs between task performance and average recurrent depth under different adaptive inference thresholds. Markers denote mean success rates and error bars indicate one standard deviation.
Figure 12 : Selected trajectory segments from a single episode in the Scene environment. FlexLoop adaptively uses fewer recurrent steps for simple, monotonic motions and increases computation at challenging decision points, such as goal switching.
Figure 13 : Depth-wise performance evolution of pretrained IRU-5 and FlexLoop from scratch on both training and unseen tasks. Solid lines denote mean success rates and shading indicates one standard deviation.
Figure 14 : Elasticity and Pareto curves of pretrained IRU-5 and FlexLoop from scratch on both training and unseen tasks. Markers denote mean success rates and the shading and error bars indicate one standard deviation.
Figure 15 : Elasticity and Pareto Curves of the policies with different pretraining recurrent depths before and after FlexLoop post-training. Markers denote mean success rates and shading indicates one standard deviation.