FlexLoop enables adaptive computation in reinforcement learning policies
FlexLoop: Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL
Machine LearningArtificial Intelligence
Summary
Deep reinforcement learning policies use repeated steps to make decisions, but often rely heavily on running all of these steps even when fewer would work. The authors found that existing methods struggle to make good decisions at fewer steps after training. They created FlexLoop, a technique that trains the policy to work reliably at various step counts, allowing it to adapt in real time and save computing resources. Experiments show FlexLoop keeps strong performance while using less computation and running faster.
What this means in practice
- •For autonomous vehicle developers: Deploy adaptive control policies that reduce computation while maintaining decision quality for long driving tasks.
- •For robotics engineers: Implement flexible robot policies that adjust computation on the fly to balance speed and reliability during extended operations.
Authors
Xun Wang, Ruishuo Chen, Yu Chen, Zhuoran Li, Longbo Huang
Abstract
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.