Paper deep dive
FlexLoop: Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL
Xun Wang, Ruishuo Chen, Yu Chen, Zhuoran Li, Longbo Huang
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2609.34488v1
- Canonical: https://arxiv.org/abs/2609.34488v1
Trouble viewing inline? Open PDF directly โ
Full Text
66,475 characters extracted from source content.
Expand or collapse full text
arXiv preprint, September 2026 FlexLoop: Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Xun Wang 1 , Ruishuo Chen 1 , Yu Chen 1 , Zhuoran Li 1 , and Longbo Huang 1# 1 Institute for Interdisciplinary Information Sciences, Tsinghua University # Correspondence: longbohuang@tsinghua.edu.cn 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 special- ization: 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 consis- tency. 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 to43% and achieves up to 1.34ร wall-clock speedup in a stress test. 1. Introduction 12345 Recurrence Depth K 0.00 0.25 0.50 0.75 1.00 Success Rate IRU-1IRU-2IRU-5 Figure 1|Fixed-depth training induces recurrent-depth specialization. Each IRU-ํพ is trained at theํพ-th recurrent depth. Re- liable decisions are concentrated near the full trained depth. Scaling model capacity has been a central driver of progress in modern machine learning, with larger models enabling increasingly capable behavior across language modeling (Ka- plan et al., 2020; Wei et al., 2022), robotics (Zitkovich et al., 2023; Kim et al., 2024), and deep reinforcement learning (DRL) (Wang et al., 2025). Beyond increasing parameter count, a complementary line of work has explored looped architectures (Dehghani et al., 2019; Yang et al., 2024; Goyal et al., 2026), which repeatedly apply a shared compu- tation block to increase effective computational depth while keeping the number of trainable parameters fixed. By scal- ing computation rather than model size, such architectures provide a parameter-efficient way to expand computational capacity. 1 arXiv:2609.34488v1 [cs.LG] 28 Sep 2026 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Recent work on looped recurrent policies, such as IRU (Ghugare et al., 2026), has shown that recurrent computation provides an effective scaling dimension for DRL, with additional recurrent steps substantially improving planning and decision making on long-horizon tasks without increasing policy parameter count. This naturally suggests recurrent depth as a deployment-time computation budget, analogous to recent looped language models that support budget-aware or adaptive inference (Jeddi et al., 2026; Popescu et al., 2026; Schwethelm et al., 2026). Ideally, such policies shall exhibit depth elasticity, namely, the ability of a single policy to maintain reliable decision quality across multiple recurrent depths and flexibly adjust its inference computation at deployment. Surprisingly, as shown in Figure 1, standard RL training does not yield such elasticity. Instead, looped policies trained with a fixed recurrent depth exhibit severe recurrent-depth specialization, with reliable decision-making concentrated at the full training depth and performance deteriorating sharply at shallower depths. This specialization ties deployment to the full recurrent depth. Although the full-depth policy itself is strong, always executing the maximum recurrent depth may use more computation than necessary for some decisions and limit adaptation to deployment resource constraints. This motivates a key question: Can we preserve the full-depth capability of a looped RL policy while making its shallower recurrent depths reliable for decision making? Addressing this question is nontrivial. Empirically, we find that looped architectures have sufficient capacity to support useful policies across multiple recurrent depths, yet naively training these depths jointly suppresses full-depth performance, while relying on distillation alone can degrade generalization to unseen tasks. To resolve this dilemma, we propose FlexLoop, a novel post-training framework that converts a pre- trained fixed-depth looped policy, which is reliable primarily at its training depth, into an elastic-depth policy that remains effective across recurrent depths. Starting from the pretrained policy, FlexLoop keeps training on the original RL objective to preserve full-depth performance and out-of-distribution generalization, while introducing adjacent-depth policy distillation to progressively propagate deci- sion quality toward shallower recurrent steps. The resulting consistency across successive depths further provides a natural signal for state-wise adaptive inference, allowing computation to terminate once the policy has sufficiently stabilized. Across 30 online and offline long-horizon goal-conditioned environments with discrete and contin- uous action spaces, FlexLoop preserves full-depth performance while making shallower recurrent depths reliable, improving the deployment-time computationโperformance trade-off. With state-wise adaptive inference, FlexLoop automatically allocates recurrent computation according to state dif- ficulty. It reduces the average recurrent depth by up to 43% and achieves up to 1.34รwall-clock speedup in a large-scale stress test, while maintaining competitive or better task performance. In summary, our main contributions are: โขWe identify recurrent-depth specialization in existing looped RL policies: despite their capacity to operate across recurrent depths, fixed-depth training fails to induce depth elasticity, concentrating reliable decisions at the training depth and limiting adaptation to deployment computation constraints. โข We propose FlexLoop, a novel post-training framework that induces depth elasticity in pretrained looped RL policies. It keeps training on the original RL objective to preserve full-depth capability and performs adjacent-depth policy distillation to transfer decision quality from deeper to shallower 2 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL recurrent depths, enabling state-wise adaptive inference at deployment without retraining. โขAcross 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth performance while enabling reliable decisions at shallower recurrent depths. State-wise adaptive inference reduces average recurrent computation by up to 43% and achieves up to 1.34ร wall-clock speedup, while maintaining competitive or better task performance. 2. Related Work Computation Scaling in DRL. Prior work has incorporated additional computation into DRL through planning and iterative value estimation (Tamar et al., 2016; Silver et al., 2017; Farquhar et al., 2018), while recent studies show that increasing network depth can substantially improve long-horizon decision making (Wang et al., 2025). More recently, Ghugare et al. (2026) showed that recurrent computation can scale RL capability without increasing policy parameter count. We build on this direction by investigating whether recurrent policies can remain effective across inference depths. We find that standard fixed-depth training instead leads to recurrent-depth specialization, highlighting the need for elastic recurrent policies. Looped Architectures and Elastic Depth. Looped architectures, which reuse the same parameters across multiple computation steps, have been studied for iterative reasoning and parameter-efficient computation (Dehghani et al., 2019; Yang et al., 2024; Zhu et al., 2025; Du et al., 2026). Recent work extends looped models to variable inference budgets. For example, LoopFormer trains language models across different recurrent depths (Jeddi et al., 2026), while ELT uses intra-loop self-distillation to transfer behavior from high-compute to intermediate depths in visual generation (Goyal et al., 2026). Others study adaptive stopping and efficient depth-adaptive inference in looped language models (Popescu et al., 2026; Schwethelm et al., 2026). However, directly extending such depth elasticity to RL is challenging. Unlike settings with explicit targets, RL provides no ground-truth intermediate decisions, as policy quality is ultimately determined by long-horizon return rather than local output agreement. Achieving reliable depth elasticity for looped RL policies therefore remains underexplored. 3. Preliminary Goal-Conditioned RL. We consider a goal-conditioned Markov decision processM=(S,A,G, ํ, ํ, ํ 0 , ํพ), whereS,A, andGdenote the state, action, and goal spaces, respectively;ํ(ยท | ํ , ํ)is the transition kernel;ํ(ํ , ํ, ํ)is the goal-conditioned reward;ํ 0 (ํ , ํ)is a joint distribution over initial states and goals; andํพ โ [0,1)is the discount factor. At the beginning of each episode, (ํ 0 , ํ) โผ ํ 0 , and a goal-conditioned policyํ(ํ | ํ , ํ)induces trajectories according toํ ํก โผ ํ(ยท | ํ ํก , ํ) and ํ ํก+1 โผ ํ(ยท | ํ ํก , ํ ํก ). The objective is to maximize ํฝ(ํ)=ํผ " ํปโ1 โ๏ธ ํก=0 ํพ ํก ํ(ํ ํก , ํ ํก , ํ) # . We consider both online and offline settings. In the latter, the agent learns from a fixed datasetD without further environment interaction. Looped Policies and IRU. Looped architecture (Jeddi et al., 2026; Zhu et al., 2025; Du et al., 2026) increase computational depth by repeatedly applying a shared computation block, without 3 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL introducing depth-specific parameters. Letํฅdenote the network input andโ ํ the representation after ํ recurrent steps. We write a looped computation abstractly as โ ํ+1 = ํน ํ (โ ํ , ํฅ), where ํน ํ is the computation block whose parameters ํ are shared across recurrent depth. Recent work (Ghugare et al., 2026) introduced the Interpolation Recurrent Unit (IRU) as a minimal instantiation of this architecture and demonstrated that scaling its recurrent depth can substantially improve long-horizon decision making in DRL. We therefore focus primarily on IRU-based policies. Given recurrent state ํ ํ , IRU computes ํ ํ = ํ ( FORGET ํ ([ํฅ, tanh(ํ ํ )]) ) , ํผ ํ = tanh ( INPUT ํ ([ํฅ, tanh(ํ ํ )]) ) , ํ ํ+1 = ํ ํ โ ํ ํ +(1โ ํ ํ )โ ํผ ํ . At each depthํ, an output head producesํง ํ = ํบ ํ (ํ ํ ), which may correspond to action values ํ ํ (ํ , ํ,ยท), a policyํ ํ (ยท | ํ , ํ), or a value estimateํ ํ (ํ , ํ)depending on the RL algorithm. We denote the full recurrent depth byํพ. While standard IRU policies are optimized throughํง ํพ , we aim to make ํง ํ reliable for ํ < ํพ while preserving full-depth capability. 4. Method: FlexLoop As shown in Figure 2, we propose FlexLoop, a novel post-training framework for converting a depth- specialized looped RL policy into an elastic one. Starting from a pretrained loop policy, FlexLoop keeps training on the original RL training while introducing adjacent-depth policy self-distillation. The RL objective preserves the task capability acquired at full depth, while depth-wise distillation progressively makes shallower recurrent depths reliable for decision making, and induces a natural policy-consistency signal for state-wise adaptive-depth inference. 4.1. Elastic Post-Training FlexLoop starts from a pretrained looped policy with full recurrent depthํพ, inheriting all pretrained model parameters. During post-training, it continues optimization with the original algorithm-specific RL objectiveL RL . We denote the policy induced after ํ recurrent steps by ํ ํ (ยท | ํ , ํ). Full-Depth RL Objective. FlexLoop inherits all parameters from the pretrained depth-specialized model and continues to optimize its full-depth value functions and policy under the same algorithm- specific RL objectiveL RL as in pretraining. This retains the capability acquired at depthํพand preserves the out-of-distribution generalization of the pretrained policy while shallower recurrent depths are elasticized. Progressive Depth-Wise Policy Distillation. To make shallower recurrent depths directly exe- cutable, FlexLoop progressively transfers decision quality from deeper to shallower recurrent policies. Let ํ ํ (ยท | ํ , ํ) denote the policy induced after ํ recurrent steps. We optimize L depth = 1 ํพโ 1 ํพโ1 โ๏ธ ํ=1 ํท ( sg[ํ ํ+1 ], ํ ํ ) ,(1) 4 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL โ ! โ " โ # ... Looped Block ํ # (ํ ,ํ) ํ ! ํ " ... PretrainedPolicy โ ! โ " โ # ... Looped Block ํ # ํ ! ํ " ... ElasticDecision ํท !" ํท !" ํท !" โ !" โ #$%&' Poor PoorGood โ ! โ " โ # ... Looped Block ํ # (ํ ,ํ) ํ ! ํ " ... Good Good Good ํ ! ํ $%&' ํ ! ํ (%)* ํ ! ํ +$*,-+ ํฟ ! ํฟ ! ํฟ ! ํฟ " ํฟ # ํ " ํ " ํ # ํ " ํ . FlexLoopPost-TrainingState-WiseAdaptiveInference (a)(b)(c) ํฟ # =ํท !" (ํ #$% ||ํ # ) 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. whereํทis an algorithm-specific policy discrepancy andsg[ยท]denotes stop-gradient. This progressive distillation propagates full-depth decision quality toward shallower recurrent steps, while explicitly controlling policy variation across adjacent depths. The resulting discrepancy between successive policies provides a natural signal for state-wise adaptive inference, allowing computation to stop once the policy has sufficiently stabilized. The full FlexLoop objective is L FlexLoop =L RL + ํ depth L depth (2) 4.2. Instantiating FlexLoop across RL Algorithms Following the setting of Ghugare et al. (2026), we instantiate FlexLoop on three representative DRL algorithms for goal-conditioned long-horizon tasks, spanning online and offline settings as well as discrete and continuous action spaces. Goal-Conditioned DQN (GC-DQN). We first instantiate FlexLoop with a goal-conditioned variant of DQN (Mnih et al., 2015). Since a discrete-action policy is induced directly by its action values, we convert each depth-ํ value output into a soft policy ํ ํ ํ (ํ | ํ , ํ)= exp ( ํ ํ (ํ , ํ, ํ)/ํ ) ร ํ โฒ exp ( ํ ํ (ํ , ํ, ํ โฒ )/ํ ) ,(3) whereํis a temperature parameter. We then apply progressive policy distillation between adjacent depths, L DQN depth = 1 ํพโ 1 ํพโ1 โ๏ธ ํ=1 ํผ (ํ ,ํ)โผD h ํท KL sg[ํ ํ ํ+1 ]โฅ ํ ํ ํ i ,(4) whereD denotes the sampled transition batch from the replay buffer. Goal-Conditioned PPO (GC-PPO). We next instantiate FlexLoop with a goal-conditioned variant of PPO (Schulman et al., 2017). We transfer decision quality across recurrent depth by distilling adjacent policies: L PPO depth = 1 ํพโ 1 ํพโ1 โ๏ธ ํ=1 ํผ (ํ ,ํ)โผB [ ํท KL ( sg[ํ ํ+1 (ยท | ํ , ํ)]โฅ ํ ํ (ยท | ํ , ํ) )] ,(5) whereB denotes the on-policy rollout batch. 5 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Goal-Conditioned IQL (GC-IQL). We finally instantiate FlexLoop with a goal-conditioned variant of IQL (Kostrikov et al., 2022), where learning is restricted to a fixed offline datasetD. We distill the Gaussian policy across adjacent recurrent steps, L IQL depth = 1 ํพโ 1 ํพโ1 โ๏ธ ํ=1 ํผ (ํ ,ํ)โผD e ํท KL ( sg[ํ ํ+1 (ยท | ํ , ํ)]โฅ ํ ํ (ยท | ํ , ํ) ) ,(6) where e ํท KL denotes the normalized policy KL used for continuous-action distillation, avoiding the scale-inflation degeneracy of the raw KL objective. 4.3. State-Wise Adaptive Inference Once a policy produces meaningful decisions at multiple recurrent depths, the remaining question is when additional computation is useful. We first observe that good endpoint performance alone does not, in general, guarantee good shallower decisions. Proposition 4.1 (No Shallow-Depth Guarantee). For any MDPMwith bounded returns, recurrent depth ํพ โฅ 2, and ํ > 0, there exists a shared-weight recurrent policyฮ such that ํฝ(ฮ ํพ ) โฅ sup ํ ํฝ(ํ)โ ํ, ํฝ(ฮ ํ ) โค inf ํ ํฝ(ํ)+ ํ, โ1โค ํ < ํพ, where sup ํ and inf ํ are taken over all stationary goal-conditioned Markov policies. Proposition 4.1 shows that strong full-depth performance alone does not imply useful intermediate decisions. The recurrent states may instead serve purely as internal computation. The proof is given in Appendix A.1. ThroughL depth , FlexLoop explicitly controls how the policy evolves across recurrent depth. Define the adjacent-depth discrepancy ํ ํ = sup ํ ,ํ ํท TV ํ ํ (ยท | ํ , ํ), ํ ํ+1 (ยท | ํ , ํ) ,(7) then we have the following proposition. Proposition 4.2 (Shallow-Depth Performance Bound). Assumeโ(ํ , ํ, ํ) โSรAรG,|ํ(ํ , ํ, ํ)| โค ํ max and ํพ < 1. Define ํฟ max = max 1โค ํ<ํพ sup (ํ ,ํ) ํท KL ํ ํ+1 (ยท | ํ , ํ)โฅ ํ ํ (ยท | ํ , ํ) . Then, for any 1โค ํ < ํพ, | ํฝ(ํ ํ )โ ํฝ(ํ ํพ ) | โค 2ํ max (1โ ํพ) 2 (ํพโ ํ) โ๏ธ ํฟ max 2 . Proposition 4.2 shows that controlling the largest adjacent-depth policy discrepancy directly bounds the performance gap between any shallow policy and the full-depth policy. Thus, when FlexLoop keeps neighboring recurrent policies close, shallow-depth performance cannot drift arbitrarily far from the full-depth endpoint. The proof is given in Appendix A.2. State-Wise Adaptive Inference. Motivated by Proposition 4.2, we monitor the depth-wise change in the policy output, ํฟ ํ (ํ , ํ)= ํท ( ํ ํโ1 (ยท | ํ , ํ), ํ ํ (ยท | ํ , ํ) ) , ํ > 1(8) where ํท is a policy-space divergence. 6 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 Average max k ฮด k Policy Discrepancy over Training 2345678910 Recurrence Depth 0.0 0.1 0.2 0.3 0.4 Average ฮด k Policy Discrepancy across Depth FlexLoop (Ours)IRU-10 (a) Reduction of adjacent-depth policy discrepancy. 10 โ3 10 โ2 10 โ1 10 0 ฮด k 0.0 0.2 0.4 0.6 0.8 KL( ฯ k โ ฯ K ) Train 10 โ3 10 โ2 10 โ1 10 0 ฮด k Unseen FlexLoop (Ours)IRU-10 (b) Local-to-endpoint policy consistency. Figure 3|(a) FlexLoop reduces adjacent-depth policy discrepancy during post-training and across depths. (b) Smallerํฟ ํ corresponds to smallerํท KL (ํ ํ โฅํ ํพ )on both training and unseen tasks. Marker size denotes the sample fraction in each pooled quantile bin. Figure 3a shows that FlexLoop reduces adjacent-depth policy divergence during training and across recurrent depths, while Figure 3b shows that smaller local discrepancies are consistently associated with smaller divergence to the full-depth policy on both training and unseen tasks. This suggests that ํฟ ํ is informative of endpoint proximity, motivating its use as a practical stopping signal for state-wise adaptive inference. Accordingly, given a threshold ํ, minimum depth ํ min , inference terminates at ํ โ (ํ , ํ)= min ( ํ > ํ min : ํฟ ํ (ํ , ํ) < ํ โชํพ ) .(9) Varying ํ yields different computationโperformance trade-offs from a single trained policy. 5. Experiments In this section, we evaluate FlexLoop across diverse online and offline goal-conditioned RL settings to answer three questions: (i) can it preserve full-depth capability while enabling reliable shallow-depth decisions? (i) how does it trade off performance and inference compute? and (i) how do training design choices and recurrent architectures affect depth elasticity? 5.1. Experimental Setup Benchmarks. Following Ghugare et al. (2026), we evaluate FlexLoop on 30 long-horizon goal- conditioned environments spanning online and offline RL, discrete and continuous action spaces, as well as in-distribution learning and out-of-distribution generalization. Specifically, we consider four environments from the BoxPick stitching benchmark (Bortkiewicz et al., 2025), where an agent moves boxes on 2-D maps to designated goal locations; one LightsOut environment (Anderson & Feil, 1998), formulated as a 4ร5 binary switch puzzle; and 25 environments from OGBench (Park et al., 2025), covering three manipulation scenarios and two navigation scenarios. For BoxPick and LightsOut, we additionally evaluate generalization to unseen instances. Detailed environment descriptions are provided in Appendix B.1. We use goal-conditioned DQN for BoxPick, PPO for LightsOut, and goal-conditioned IQL for OGBench as the underlying RL algorithms. Baselines and Implementation. We use IRU as the common looped-policy backbone and denote a pretrained depth-ํพIRU policy as IRU-ํพ. Following the results reported by Ghugare et al. (2026), we chooseํพ=5 as the default setting, as it provides strong performance while maintaining a practical computational cost. FlexLoop is initialized from the corresponding IRU-5 checkpoint and post-trained for a fixed number of optimization steps to acquire depth elasticity. Our implementation builds on the codebase of Ghugare et al. (2026), retaining the same environment configurations, RL 7 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 0.00.20.4 Env Steps (B) 0.0 0.5 1.0 Success Rate IRU-5 Pretraining 0.00.20.4 Env Steps (B) FlexLoop (Ours) BoxPick k=1k=2k=3k=4k=5 0.000.250.500.751.00 Gradient Steps (M) 0.0 0.2 0.4 Success Rate IRU-5 Pretraining 0.000.250.500.751.00 Gradient Steps (M) FlexLoop (Ours) OGBench k=1k=2k=3k=4k=5 012 Env Steps (B) 0.0 0.5 Success Rate IRU-5 Pretraining 012 Env Steps (B) FlexLoop (Ours) LightsOut 4x5 k=1k=2k=3k=4k=5 (a) Depth-wise decision performance on training tasks. 0.00.20.4 Env Steps (B) 0.00 0.25 0.50 0.75 Success Rate IRU-5 Pretraining 0.00.20.4 Env Steps (B) FlexLoop (Ours) BoxPick k=1k=2k=3k=4k=5 012 Env Steps (B) 0.00 0.05 0.10 Success Rate IRU-5 Pretraining 012 Env Steps (B) FlexLoop (Ours) LightsOut 4x5 k=1k=2k=3k=4k=5 (b) Depth-wise decision performance on unseen tasks. 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. algorithms, and evaluation protocols unless otherwise stated. We report results over three random seeds. Additional implementation details and hyperparameters are provided in Appendix B.2. 5.2. Analysis of Learned Depth Elasticity We evaluate FlexLoop on 30 long-horizon goal-conditioned environments using pretrained IRU-5 policies. Figure 4 tracks the evolution of the recurrent-depth performance during IRU pretraining and FlexLoop post-training, showing the emergence of reliable shallow-depth policies and their generalization to unseen tasks in the BoxPick and LightsOut environments. Full environment-wise curves are provided in Appendix C.1. We further compare the depth-wise performance curves of the final pretrained and post-trained policies in Appendix C.2, providing a direct characterization of their depth elasticity. These results reveal several key observations: FlexLoop improves shallow-depth decisions while preserving full-depth capability. Across almost all tasks in Figure 4a, pretrained IRU policies concentrate reliable decision performance at the full depth, whereas shallower recurrent depths remain ineffective. After post-training, FlexLoop maintains or further improves the full-depth performance while rapidly increasing the performance of shallow recurrent depths. This shows that FlexLoop successfully converts a depth-specialized policy into an elastic policy without sacrificing the capability acquired by the pretrained model. FlexLoop generalizes depth elasticity beyond the training distribution. We evaluate the post- trained policies on unseen tasks to examine whether the learned depth elasticity generalizes beyond the training environments. As shown in Figure 4b, FlexLoop improves shallow-depth performance on unseen tasks while preserving the full-depth capability of the pretrained policy. This suggests that the improved shallow policies retain the general decision capability of the pretrained policy rather than simply overfitting to the training environments. Shallower recurrent depths are more challenging to elasticize. In Figure 4, the improvement speed during post-training decreases monotonically with shallower recurrent depths. A possible explanation is that a reliable elastic policy needs to simultaneously support effective standalone decision making at shallow depths, maintain consistency across recurrent depths, and preserve robustness under distribution shifts. With fewer recurrent steps available, shallow policies have less 8 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL computation capacity to refine intermediate representations and satisfy these requirements, making the shallowest depths harder to elasticize. The underlying RL formulation influences how FlexLoop acquires depth elasticity. FlexLoop achieves the largest improvements on the GC-DQN-based BoxPick benchmark, while gains are smaller on PPO and offline IQL tasks. We hypothesize that this difference arises from the coupling between value learning and policy execution. In value-based methods, improved value estimation directly determines action selection across recurrent depths. In contrast, methods with separately parameterized policies, such as PPO and IQL, rely on indirect policy updates from value learning, making depth-wise adaptation more challenging. 5.3. State-wise Adaptive Inference FlexLoopโs learned depth elasticity enables adaptive recurrent computation during deployment. We first evaluate whether adaptive inference can reduce computation while preserving performance, and then analyze the computation-performance trade-offs induced by different stopping thresholds. Table 1|State-wise adaptive inference results across en- vironments. ยฑ indicates one standard deviation. Environment FlexLoopIRU-5 Avg. DepthSucc.Succ. OGBench4.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 Adaptive Inference Performance. By tun- ing the stopping thresholdํin Equation (9), we select adaptive inference policies for each environment. Table 1 summarizes the result- ing performance and average recurrent depth. FlexLoop maintains the performance of the full-depth policy while reducing the average recurrent depth by up to 43%, demonstrating that learned depth elasticity enables effective inference computation reduction without sac- rificing performance. Full environment-wise results are provided in Appendix C.3. Computation-Performance Trade-off. We further analyze how the stopping threshold controls the trade-off between computation and performance. By varyingํ, we obtain the Pareto curves ( See Figure 11 in Appendix C.4). In general, smaller thresholds allocate deeper recurrent computation and tend to achieve higher performance, while larger thresholds reduce the average inference depth with controllable performance trade-offs. The resulting Pareto curves show that FlexLoop often matches Table 2|Wall-clock comparison on BoxPick-Exact-4.ยฑ indicates one standard deviation. Bold numbers indicate the shortest wall-clock latency. ModelSuccessAvg. Depth Wall-clock (ms) Train IRU-50.995ยฑ 0.0015170.1ยฑ 17.2 FlexLoop 0.993ยฑ 0.001 3.85ยฑ 0.04133.0ยฑ 7.7 IRU-10 0.997ยฑ 0.00010333.2ยฑ 36.5 FlexLoop 0.996ยฑ 0.001 9.35ยฑ 0.04 319.9ยฑ 22.2 Unseen IRU-50.511ยฑ 0.0855170.1ยฑ 17.2 FlexLoop 0.543ยฑ 0.093 3.80ยฑ 0.05127.4ยฑ 6.3 IRU-10 0.664ยฑ 0.08010333.2ยฑ 36.5 FlexLoop 0.675ยฑ 0.066 9.37ยฑ 0.04 321.1ยฑ 22.7 the pretrained IRU policy performance with reduced average recurrent depth. In addition, this favorable trade-off is preserved on unseen tasks, demonstrating that FlexLoopโs learned depth elasticity enables generalizable adaptive computation beyond the training distribution. Wall-Clock Time Saving. To evaluate the practical inference speedup enabled by adap- tive recurrent depth, we additionally train an IRU-10 policy and apply FlexLoop post- training to obtain the corresponding elastic variant. We evaluate both IRU-5 and IRU-10 fixed-depth policies and their FlexLoop variants on BoxPick-Exact-4 with a batch of 32,768 parallel environments. As shown in Table 2, FlexLoop 9 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL maintains competitive performance relative to the pretrained policies while using fewer recurrent steps, achieving up to 1.34ร wall-clock speedup compared with fixed-depth inference. 5.4. Ablation Study 12345 Recurrence Depth 0.00 0.25 0.50 0.75 1.00 Success Rate Train 12345 Recurrence Depth Unseen FlexLoop (Ours)Distill OnlySelf-TD OnlyAnyK-TDIRU-5 (a) Component Ablation. 12345 Recurrence Depth 0.00 0.25 0.50 0.75 1.00 Success Rate Train 12345 Recurrence Depth Unseen ฮป depth = 0 ฮป depth = 0.25 ฮป depth = 0.5 ฮป depth = 1 ฮป depth = 2 ฮป depth = 3 IRU-5 (b) Sensitivity to ํ depth . 12345 Recurrence Depth 0.00 0.25 0.50 0.75 1.00 Success Rate Train 12345 Recurrence Depth Unseen IRULSTM-256LSTM-184FlexLoopPretrain (c) Architectural Generality. Figure 5|Ablation studies of FlexLoop. Solid lines denote mean performance, while shaded regions indicate one standard deviation. Based on BoxPick-Exact-4, we study FlexLoopโs key components,ํ depth sensitivity, and architectural generality. Additional results on from-scratch training and pretraining recurrent depth are provided in Appendix D. Contribution of Key Components. We ablate the key components of FlexLoop, including Distill Only, which performs depth-wise distillation from the pretrained full-depth policy, Self-TD Only, which solely keeps the original RL training, and AnyK-TD, which jointly optimizes all recurrent depths with TD learning. As shown in Figure 5a, Self-TD Only preserves the pretrained policyโs generalization but provides limited shallow-depth improvement, showing that depth elasticity does not emerge solely from the original RL objective. Meanwhile, AnyK-TD and Distill Only each suffer from limitations in full-depth capability or unseen-task generalization. These highlight the distinct roles of the two objectives. Continued RL optimization preserves full-depth capability and generalization, while depth-wise policy distillation transfers decision quality to shallower depths. Sensitivity to ํ depth . We further study the effect of the depth distillation weight ํ depth . As shown in Figure 5b, FlexLoop remains effective across a broad range ofํ depth values. Increasingํ depth strengthens depth-wise distillation and consistently improves shallow-depth performance on training tasks. However, excessive distillation can hurt generalization to unseen tasks, where a moderate ํ depth achieves a better balance between shallow-depth adaptation and generalization. Generality across Architectures. We further apply FlexLoop to LSTM-based policies (Hochreiter & Schmidhuber, 1997) to evaluate its architectural generality. LSTM-256 and LSTM-184 denote backbones with hidden sizes of 256 and 184, respectively, where LSTM-184 has a comparable parameter count to IRU. As shown in Figure 5c, both LSTM variants exhibit depth specialization during pretraining, while FlexLoop consistently improves shallow-depth performance and preserves full-depth capability across architectures. Moreover, LSTM-256 benefits more from FlexLoop than LSTM-184, suggesting that larger capacity allows a policy to acquire stronger depth elasticity. 5.5. Visualization of Adaptive Inference To qualitatively understand how FlexLoop utilizes its learned depth elasticity, we visualize the recurrent computation allocated during successful trajectories on three OGBench tasks. Figure 6 shows the recurrent depth selected by FlexLoop at each environment step, with representative 10 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 050100150200 Environment Step 1 2 3 4 5 Executed Depth K (a) Scene. 2 3 4 5 Executed Depth K StartGoal (b) AntMaze. 2 3 4 5 Executed Depth K StartGoal (c) PointMaze. Figure 6|State-wise adaptive recurrent-depth allocation across OGBench tasks. FlexLoop allocates more recurrent computation to states requiring more complex decisions. trajectory segments visualized frame by frame in Appendix C.5. The allocated computation appears to track state-level decision difficulty. In Scene, FlexLoop increases recurrent depth at skill boundaries, such as goal switching, while using only three or two recurrent steps during simpler execution phases such as button pressing. Similarly, deeper computation is concentrated near turns and goals in AntMaze, whereas straight navigation uses shallower inference. On the simpler PointMaze task, FlexLoop consistently terminates after two recurrent steps. These trajectories suggest that FlexLoop reallocates recurrent computation toward challenging decisions and terminates earlier once adjacent- depth policies have stabilized, rather than uniformly reducing computation. 6. Conclusion and Future Directions In this work, we identify recurrent-depth specialization in pretrained looped RL policies and propose FlexLoop, a post-training framework that induces depth elasticity through continued RL optimization and adjacent-depth policy distillation. Across 30 online and offline long-horizon goal-conditioned environments, FlexLoop preserves full-depth capability while enabling reliable shallower-depth inference, reducing average recurrent computation by up to 43% and achieving up to 1.34รwall- clock speedup, while maintaining competitive or better task performance. Future work includes extending depth elasticity beyond the trained depth, developing learning-based adaptive inference mechanisms beyond the current depth-consistency criterion, and studying depth elasticity in larger- scale looped architectures such as transformer-based models. References Anderson, M. and Feil, T. Turning lights out with linear algebra. Mathematics magazine, 71(4):300โ303, 1998. Bortkiewicz, M., Paลucki, W., Ostaszewski, M., and Eysenbach, B. Is temporal difference learning the gold standard for stitching in rl?, 2025. URL https://arxiv.org/abs/2510.21995. Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, L. Universal transformers. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=HyzdRiR9Y7. Du, H., Dong, Y., and Ning, X. Latent thinking optimization: Your latent reasoning language model secretly encodes reward signals in its latent thoughts. In Vondrick, C., Hariharan, B., Raffel, C., Pinto, L., Yang, D., and Faust, A. (eds.), International Conference on Learning Representations, volume 2026, p. 101000โ101032, 2026. URL https://proceedings.iclr.c/paper_files/paper/2026/file/ a4b1a25e1b37a17733c1fea50d0b7446-Paper-Conference.pdf. Farquhar, G., Rocktaeschel, T., Igl, M., and Whiteson, S. TreeQN and ATreec: Differentiable tree planning for deep reinforcement learning. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=H1dh6Ax0Z. 11 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Ghugare, R., Bortkiewicz, M., Ziarko, A., and Eysenbach, B. On the role of computation in reinforcement learning. In Forty-third International Conference on Machine Learning, 2026. URL https://openreview.net/forum?id=HKzMnBX5B2. Goyal, S., Agrawal, S., Anil, G. G., Jain, P., Paul, S., and Kusupati, A. Elt: Elastic looped transformers for visual generation. In Favaro, P., Kukelova, Z., Maki, A., Rohrbach, A., Schindler, K., and Tombari, F. (eds.), Computer Vision โ ECCV 2026, p. 308โ326, Cham, 2026. Springer Nature Switzerland. ISBN 978-3-032-37324-3. Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural Computation, 9(8):1735โ1780, 1997. doi: 10.1162/neco.1997.9.8.1735. Jeddi, A., Ciccone, M., and Taati, B. Loopformer: Elastic-depth looped transformers for latent reasoning via shortcut modulation. In The Fourteenth International Conference on Learning Representations, 2026. URL https://openreview.net/forum?id=RzYXb5YWBs. Kakade, S. and Langford, J. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning, ICML โ02, p. 267โ274, San Francisco, CA, USA, 2002. Morgan Kaufmann Publishers Inc. ISBN 1558608737. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models, 2020. URL https://arxiv.org/abs/2001.08361. Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Balakrishna, A., Nair, S., Rafailov, R., Foster, E. P., Sanketi, P. R., Vuong, Q., Kollar, T., Burchfiel, B., Tedrake, R., Sadigh, D., Levine, S., Liang, P., and Finn, C. OpenVLA: An open-source vision-language-action model. In 8th Annual Conference on Robot Learning, 2024. URL https://openreview.net/forum?id=ZMnD6QZAE6. Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=68n2s9ZJWF8. Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature, 518 (7540):529โ533, 2015. Park, S., Frans, K., Eysenbach, B., and Levine, S. Ogbench: Benchmarking offline goal-conditioned rl. In Yue, Y., Garg, A., Peng, N., Sha, F., and Yu, R. (eds.), International Conference on Learning Representations, volume 2025, p. 94937โ94982, 2025. URL https://proceedings.iclr.c/paper_files/paper/2025/file/ ecd92623ac899357312a8915853699-Paper-Conference.pdf. Popescu, A. C., de Ocรกriz Borde, H. S., and Liรฒ, P. Adaptive depth in looped transformers: Diagnosing learned halting gates and trajectory readouts, 2026. URL https://arxiv.org/abs/2607.20519. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Schwethelm, K., Rueckert, D., and Kaissis, G. Depth-adaptive inference of looped language models via continuous depth batching, 2026. URL https://arxiv.org/abs/2608.09444. Silver, D., van Hasselt, H., Hessel, M., Schaul, T., Guez, A., Harley, T., Dulac-Arnold, G., Reichert, D., Rabinowitz, N., Barreto, A., and Degris, T. The predictron: End-to-end learning and planning. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, p. 3191โ3199. PMLR, 06โ11 Aug 2017. URL https://proceedings.mlr.press/v70/silver17a.html. Tamar, A., WU, Y., Thomas, G., Levine, S., and Abbeel, P. Value iteration networks. In Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL https: //proceedings.neurips.c/paper_files/paper/2016/file/c21002f464c5fc5bee3b98ced83963b8-Paper.pdf. 12 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Wang, K., Javali, I., Bortkiewicz, M., Trzcinski, T., and Eysenbach, B. 1000 layer networks for self-supervised rl: Scaling depth can enable new goal-reaching capabilities. In Belgrave, D., Zhang, C., Lin, H., Pascanu, R., Koniusz, P., Ghassemi, M., and Chen, N. (eds.), Advances in Neural Information Processing Systems, volume 38, Main Conference, p. 157643โ157670. Curran Associates, Inc., 2025. doi: 10.52202/085713-5268. URL https://proceedings.neurips.c/paper_files/paper/2025/file/ e74e34c0f2d0780f34e77d8fba25b-Paper-Conference.pdf. Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., Chi, E. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https://openreview.net/forum?id=yzkSU5zdwD. Survey Certification. Yang, L., Lee, K., Nowak, R., and Papailiopoulos, D. Looped transformers are better at learning learning algorithms. In Kim, B., Yue, Y., Chaudhuri, S., Fragkiadaki, K., Khan, M., and Sun, Y. (eds.), International Conference on Learning Representations, volume 2024, p. 42195โ42214, 2024. URL https://proceedings. iclr.c/paper_files/paper/2024/file/b8402301e7f06bdc97a31bfaa653dc32-Paper-Conference.pdf. Zhu, R.-J., Wang, Z., Hua, K., Zhang, T., Li, Z., Que, H., Wei, B., Wen, Z., Yin, F., Xing, H., Li, L., Shi, J., Ma, K., Li, S., Kergan, T., Smith, A., Qu, X., Hui, M., Wu, B., Min, Q., Huang, H., Zhou, X., Ye, W., Liu, J., Yang, J., Shi, Y., Lin, C., Zhao, E., Cai, T., Zhang, G., Huang, W., Bengio, Y., and Eshraghian, J. Scaling latent reasoning via looped language models, 2025. URL https://arxiv.org/abs/2510.25741. Zitkovich, B., Yu, T., Xu, S., Xu, P., Xiao, T., Xia, F., Wu, J., Wohlhart, P., Welker, S., Wahid, A., Vuong, Q., Vanhoucke, V., Tran, H., Soricut, R., Singh, A., Singh, J., Sermanet, P., Sanketi, P. R., Salazar, G., Ryoo, M. S., Reymann, K., Rao, K., Pertsch, K., Mordatch, I., Michalewski, H., Lu, Y., Levine, S., Lee, L., Lee, T.-W. E., Leal, I., Kuang, Y., Kalashnikov, D., Julian, R., Joshi, N. J., Irpan, A., brian ichter, Hsu, J., Herzog, A., Hausman, K., Gopalakrishnan, K., Fu, C., Florence, P., Finn, C., Dubey, K. A., Driess, D., Ding, T., Choromanski, K. M., Chen, X., Chebotar, Y., Carbajal, J., Brown, N., Brohan, A., Arenas, M. G., and Han, K. RT-2: Vision-language-action models transfer web knowledge to robotic control. In 7th Annual Conference on Robot Learning, 2023. URL https://openreview.net/forum?id=XMQgwiJ7KSX. 13 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL Appendix A Proofs of the Propositions15 A.1 Proof of Proposition 4.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 A.2 Proof of Proposition 4.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 B Experiment Details16 B.1 Environments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 B.2 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17 C Extended Results18 C.1 Full training curves for Figure 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 C.2 Elasticity Curves . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 C.3 Full results for Table 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 C.4 Pareto Curve . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 C.5 Visualized Trajectory Segments of Scene . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 D Additional Ablations24 D.1 FlexLoop from Scratch . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 D.2 FlexLoop on Policies with Different Pretraining Depths . . . . . . . . . . . . . . . . . . . 25 14 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL A. Proofs of the Propositions A.1. Proof of Proposition 4.1 Proof. Since the return is bounded,sup ํ ํฝ(ํ)andinf ํ ํฝ(ํ)are both finite, wheresup ํ andinf ํ are taken over all stationary goal-conditioned Markov policies. Hence, for anyํ >0 there exist policies ํ + and ํ โ such that ํฝ(ํ + ) โฅ sup ํ ํฝ(ํ)โ ํ, ํฝ(ํ โ ) โค inf ํ ํฝ(ํ)+ ํ.(10) We construct a shared-weight recurrent policyฮ =(ํน,ํบ), whereํนdenotes the shared recurrent block andํบthe policy output head. Let the recurrent state contain a scalarโ ํ , initialized asโ 0 =0, and define โ ํ+1 = ํน(โ ํ , ํ , ํ)= โ ํ + 1,(11) where ํน simply ignores (ํ , ํ), yielding โ ํ = ํ at recurrent depth ํ. Define the shared policy output head as ํบ(ยท | ํ , ํ, โ)= ( ํ โ (ยท | ํ , ํ), โ < ํพ, ํ + (ยท | ํ , ํ), โโฅ ํพ. (12) The depth-ํ policy induced byฮ is ฮ ํ (ยท | ํ , ํ) := ํบ(ยท | ํ , ํ, โ ํ ).(13) Therefore, the policy obtained after ํ recurrent steps is ฮ ํ = ํ โ , โํ < ํพ, while at the full recurrent depth, ฮ ํพ = ํ + . Using Equation (10), ํฝ(ฮ ํพ )= ํฝ(ํ + ) โฅ sup ํ ํฝ(ํ)โ ํ, and, for every ํ < ํพ, ํฝ(ฮ ํ )= ํฝ(ํ โ ) โค inf ํ ํฝ(ํ)+ ํ. Thus, a shared-weight recurrent policy can be arbitrarily close to optimal at its full depth while all of its shallower policies are arbitrarily close to the minimum achievable return.โก A.2. Proof of Proposition 4.2 Proof. For analysis, we use the equivalent discounted infinite-horizon formulation with a zero-reward absorbing terminal state, augmenting the state with time when needed, and writeํฅ= (ํ , ํ)for compactness. Define the worst-case total-variation discrepancy between adjacent depths as ํผ ํ = sup ํฅ ํท TV ํ ํ (ยท | ํฅ), ํ ํ+1 (ยท | ํฅ) , 1โค ํ < ํพ. By Pinskerโs inequality, ํผ ํ โค sup ํฅ โ๏ธ 1 2 ํท KL ํ ํ+1 (ยท | ํฅ)โฅ ํ ํ (ยท | ํฅ) โค โ๏ธ ํฟ max 2 . 15 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL For any 1โค ํ < ํพ, the triangle inequality across recurrent depths gives ํท TV ( ํ ํ (ยท | ํฅ), ํ ํพ (ยท | ํฅ) ) โค ํพโ1 โ๏ธ ํ=ํ ํผ ํ โค (ํพโ ํ) โ๏ธ ํฟ max 2 .(14) We next relate this policy discrepancy to return. By the performance difference lemma (Kakade & Langford, 2002), for any two policies ํ and ํ โฒ , ํฝ(ํ)โ ํฝ(ํ โฒ )= 1 1โ ํพ ํผ ํฅโผํ ํ ํผ ํโผํ(ยท|ํฅ) ํด ํ โฒ (ํฅ, ํ) ,(15) whereํ ํ denotes the normalized discounted stateโgoal occupancy measure underํ, andํด ํ โฒ is the advantage of ํ โฒ . Taking ํ= ํ ํ and ํ โฒ = ํ ํพ , and usingํผ ํโผํ ํพ (ยท|ํฅ) ํด ํ ํพ (ํฅ, ํ)= 0, we have ํผ ํโผํ ํ (ยท|ํฅ) ํด ํ ํพ (ํฅ, ํ) = ํผ ํโผํ ํ (ยท|ํฅ) ํด ํ ํพ (ํฅ, ํ)โํผ ํโผํ ํพ (ยท|ํฅ) ํด ํ ํพ (ํฅ, ํ) = ํผ ํโผํ ํ (ยท|ํฅ) ํ ํ ํพ (ํฅ, ํ)โํผ ํโผํ ํพ (ยท|ํฅ) ํ ํ ํพ (ํฅ, ํ) . Since |ํ| โค ํ max , ํพ < 1, |ํ ํ ํพ (ํฅ, ํ)| โค ํ max 1โ ํพ , and hence ํผ ํโผํ ํ (ยท|ํฅ) ํ ํ ํพ (ํฅ, ํ)โํผ ํโผํ ํพ (ยท|ํฅ) ํ ํ ํพ (ํฅ, ํ) โค 2ํ max 1โ ํพ ํท TV ( ํ ํ (ยท | ํฅ), ํ ํพ (ยท | ํฅ) ) . Therefore, ํผ ํโผํ ํ (ยท|ํฅ) ํด ํ ํพ (ํฅ, ํ) โค 2ํ max 1โ ํพ ํท TV ( ํ ํ (ยท | ํฅ), ํ ํพ (ยท | ํฅ) ) .(16) Substituting this into Equation (15) yields | ํฝ(ํ ํ )โ ํฝ(ํ ํพ ) | โค 2ํ max (1โ ํพ) 2 sup ํฅ ํท TV ( ํ ํ (ยท | ํฅ), ํ ํพ (ยท | ํฅ) ) .(17) Combining with Equation (14) gives | ํฝ(ํ ํ )โ ํฝ(ํ ํพ ) | โค 2ํ max (1โ ํพ) 2 (ํพโ ํ) โ๏ธ ํฟ max 2 , which proves the claim.โก B. Experiment Details B.1. Environments Following Ghugare et al. (2026), we evaluate FlexLoop on 30 long-horizon goal-conditioned environ- ments from BoxPick (Bortkiewicz et al., 2025), LightsOut (Anderson & Feil, 1998), and OGBench (Park et al., 2025). 16 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL BoxPick. BoxPick is a discrete-action benchmark, where an agent navigates in a 6ร6 grid, picks up boxes, and places them at designated goal locations. An episode is successful only when all boxes are correctly placed. We consider two types of configurations. (i) BoxPick-Exact-ํrequires the agent to transportํboxes. The training tasks involve moving boxes between two adjacent quadrants of the map, while the unseen tasks require transportation between diagonal quadrants. We evaluate ํ โ 3,4. (i) BoxPick-Generalize-ํ-ํrequires transportingํboxes in total, where training tasks only require movingํboxes and unseen tasks require moving allํboxes to their target locations. We evaluate(ํ, ํ) โ (4,1),(4,2). These four environments provide a challenging test of long-horizon planning and generalization. We use goal-conditioned DQN as the underlying RL algorithm for this environment. LightsOut. LightsOut-4x5 is a large discrete-action decision-making benchmark with a 4ร5 grid of 20 binary switches. At each step, the agent selects one switch, which toggles its own state and the states of all adjacent switches. An episode is successful when all switches are turned on. The training tasks consist of initial configurations that can be solved with relatively fewer steps, while the unseen tasks contain configurations requiring longer solution sequences. We use goal-conditioned PPO as the underlying RL algorithm for this environment. Table 3|Adjacent-depth policy distillation weight used for each environment. Environmentํ depth BoxPick-Exact-31.0 BoxPick-Exact-41.0 BoxPick-Generalize-4-1 1.0 BoxPick-Generalize-4-2 1.0 LightsOut0.05 Scene0.05 Cube0.02 Puzzle0.2 AntMaze0.03 PointMaze0.3 OGBench. OGBench is a continuous-action offline RL benchmark, where agents learn to achieve diverse test goals solely from static offline datasets. We evaluate five scenarios, including two navigation scenarios, AntMaze-Medium-Stitch and PointMaze-Large-Stitch, which test the ability to compose experiences through trajectory stitching, and three manipula- tion scenarios, Scene-Play, Cube-Double-Play, and Puzzle-4x4- Play, which evaluate fine-grained control and compositional generalization. Following the evaluation protocol of Ghugare et al. (2026), we report the average success rate over five task environments for each scenario. We use goal-conditioned IQL as the underlying RL algorithm for these environments. B.2. Implementation Details We build our implementation upon the official codebase by Ghugare et al. (2026) and reproduce the IRU baseline before implementing FlexLoop. Unless otherwise stated, we use the same training steps, environment configurations, and hyperparameters as reported in the IRU paper. For BoxPick, the temperatureํused in Equation (3) is set to 0.25. For OGBench, we use a hidden size of 256 instead of the hidden size of 64 reported in the original paper, as the former produces results closer to those reported for the benchmark. For on-policy GC-PPO, we additionally replace the GAE-based value targets used in the value loss with value estimates from the pretrained IRU model to stabilize FlexLoop training. The policy optimization objective remains unchanged and continues to use standard on-policy PPO updates. The key hyperparameter of FlexLoop is the adjacent-depth policy distillation weightํ depth . The selected values for each environment are reported in Table 3. 17 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL C. Extended Results In this section, we provide detailed environment-wise evaluation results, complementing the aggre- gated results reported in the main paper. C.1. Full training curves for Figure 4 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Exact-3 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Exact-4 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Generalize-4-1 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Generalize-4-2 k=1k=2k=3k=4k=5 0.00.51.01.52.0 Env Steps (B) 0.00 0.02 0.04 0.06 0.08 0.10 Success Rate IRU-5 Pretraining 0.00.51.01.52.0 Env Steps (B) FlexLoop (Ours) LightsOut 4x5 k=1k=2k=3k=4k=5 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. 18 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Exact-3 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Exact-4 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Generalize-4-1 k=1k=2k=3k=4k=5 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop (Ours) BoxPick-Generalize-4-2 k=1k=2k=3k=4k=5 0.00.20.40.60.81.0 Gradient Steps (M) 0.0 0.1 0.2 0.3 0.4 0.5 0.6 Success Rate IRU-5 Pretraining 0.00.20.40.60.81.0 Gradient Steps (M) FlexLoop (Ours) Scene k=1k=2k=3k=4k=5 0.00.20.40.60.81.0 Gradient Steps (M) 0.0 0.2 0.4 0.6 0.8 Success Rate IRU-5 Pretraining 0.00.20.40.60.81.0 Gradient Steps (M) FlexLoop (Ours) Cube-Double k=1k=2k=3k=4k=5 0.00.20.40.60.81.0 Gradient Steps (M) 0.0 0.1 0.2 0.3 0.4 Success Rate IRU-5 Pretraining 0.00.20.40.60.81.0 Gradient Steps (M) FlexLoop (Ours) Puzzle-4x4 k=1k=2k=3k=4k=5 0.00.20.40.60.81.0 Gradient Steps (M) 0.0 0.1 0.2 0.3 0.4 Success Rate IRU-5 Pretraining 0.00.20.40.60.81.0 Gradient Steps (M) FlexLoop (Ours) AntMaze-Medium k=1k=2k=3k=4k=5 0.00.20.40.60.81.0 Gradient Steps (M) 0.0 0.1 0.2 0.3 0.4 0.5 Success Rate IRU-5 Pretraining 0.00.20.40.60.81.0 Gradient Steps (M) FlexLoop (Ours) PointMaze-Large k=1k=2k=3k=4k=5 0.00.51.01.52.0 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 Success Rate IRU-5 Pretraining 0.00.51.01.52.0 Env Steps (B) FlexLoop (Ours) LightsOut 4x5 k=1k=2k=3k=4k=5 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. 19 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL C.2. Elasticity Curves We evaluate the depth-wise decision performance of the pretrained IRU-5 checkpoint and the cor- responding FlexLoop post-trained checkpoint at each recurrent depth, and visualize the resulting depth-performance curves for comparison. Figure 9 presents the curves on training tasks, while Figure 10 presents the corresponding results on unseen tasks. These curves provide a direct charac- terization of the depth elasticity learned by FlexLoop. 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-3 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-4 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-1 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-2 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.1 0.2 0.3 0.4 0.5 0.6 Success Rate Scene FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 Success Rate Cube-Double FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.1 0.2 0.3 0.4 Success Rate Puzzle-4x4 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.1 0.2 0.3 0.4 Success Rate AntMaze-Medium FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.1 0.2 0.3 0.4 0.5 Success Rate PointMaze-Large FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 Success Rate LightsOut 4x5 FlexLoop (Ours)IRU-5 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. 20 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-3 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-4 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-1 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-2 FlexLoop (Ours)IRU-5 12345 Recurrence Depth 0.00 0.02 0.04 0.06 0.08 0.10 Success Rate LightsOut 4x5 FlexLoop (Ours)IRU-5 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. 21 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL C.3. Full results for Table 1 Table 4 | Full state-wise adaptive inference results on all environments. ยฑ indicates one standard deviation. Task FlexLoopIRU-5 Avg. DepthSucc.Succ. Scene4.66ยฑ 0.06 0.496ยฑ 0.064 0.489ยฑ0.087 Cube4.75ยฑ 0.03 0.569ยฑ 0.056 0.492ยฑ0.114 Puzzle3.45ยฑ 0.03 0.296ยฑ 0.029 0.293ยฑ0.01 AntMaze4.80ยฑ 0.03 0.292ยฑ 0.059 0.253ยฑ0.034 PointMaze3.28ยฑ 0.30 0.395ยฑ 0.004 0.396ยฑ0.003 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 BoxPick-Exact-3(Train)4.47ยฑ 0.02 0.992ยฑ 0.004 0.992ยฑ 0.003 BoxPick-Exact-3(Unseen) 4.45ยฑ 0.05 0.768ยฑ 0.101 0.766ยฑ 0.091 BoxPick-Exact-4(Train)4.21ยฑ 0.04 0.993ยฑ 0.002 0.994ยฑ0.003 BoxPick-Exact-4(Unseen) 4.25ยฑ 0.04 0.567ยฑ 0.098 0.506ยฑ0.077 BoxPick-Gen-4-1(Train)4.68ยฑ 0.03 0.982ยฑ 0.007 0.98ยฑ 0.007 BoxPick-Gen-4-1(Unseen) 4.57ยฑ 0.02 0.82ยฑ 0.0330.82ยฑ0.022 BoxPick-Gen-4-2(Train)4.66ยฑ 0.01 0.991ยฑ 0.004 0.989ยฑ0.004 BoxPick-Gen-4-2(Unseen) 4.57ยฑ 0.01 0.956ยฑ 0.017 0.948ยฑ0.019 -Gen means -Generalize 22 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL C.4. Pareto Curve Each Pareto curve is obtained by sweeping the stopping thresholdํin Equation (9), where each threshold corresponds to one computationโperformance point. Theํgrids are selected accord- ing to the characteristics of different environment families. For BoxPick, we use a 7-point grid 0,0.005,0.01,0.02,0.03,0.05,0.1. For LightsOut, we use an 11-point uniformly spaced grid over [0,0.05]. For OGBench, we use an 11-point uniformly spaced grid for all environments: manipulation tasks (Scene, Cube, and Puzzle) use a finer grid over[0,0.005], while navigation tasks (AntMaze and PointMaze) use a grid over [0, 0.05]. 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-3 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-4 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-1 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-2 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.1 0.2 0.3 0.4 0.5 0.6 Success Rate Scene FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 Success Rate Cube-Double FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.1 0.2 0.3 0.4 Success Rate Puzzle-4x4 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.1 0.2 0.3 0.4 Success Rate AntMaze-Medium FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.1 0.2 0.3 0.4 0.5 Success Rate PointMaze-Large FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate LightsOut 4x5 FlexLoop (Adaptive)IRU-5 (Fixed K) (a) Pareto trade-offs on training tasks. 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-3 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-4 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Generalize-4-1 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Generalize-4-2 FlexLoop (Adaptive)IRU-5 (Fixed K) 12345 Average Recurrent Depth 0.000 0.025 0.050 0.075 0.100 0.125 0.150 Success Rate LightsOut 4x5 FlexLoop (Adaptive)IRU-5 (Fixed K) (b) Pareto trade-offs on unseen tasks. 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. 23 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL C.5. Visualized Trajectory Segments of Scene Figure 12 visualizes selected trajectory segments of FlexLoop in the Scene environment to illustrate its state-wise adaptive inference behavior. FlexLoop uses shallower recurrent depths for relatively simple and monotonic actions, such as lowering the robotic arm to press a button or rotating the gripper orientation, while allocating deeper computation to more critical decision points, such as switching goals after pressing the button and moving toward the drawer. t18 K=3t19 K=3t20 K=2t21 K=2t22 K=2t23 K=2t24 K=2t25 K=2 (a) Lowering the arm to press the button. t25 K=2t26 K=2t27 K=3t28 K=4t29 K=5t30 K=2t31 K=2t32 K=2 (b) Transitioning from button pressing to moving toward the drawer. t166 K=3t167 K=3t168 K=2t169 K=2t170 K=2t171 K=2t172 K=2t173 K=2 (c) Rotating the gripper orientation. 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. D. Additional Ablations D.1. FlexLoop from Scratch In this section, we further consider a more restrictive setting where the pretrained looped policy weights are unavailable and only its outputs can be queried. We ask whether FlexLoop can still acquire depth elasticity from random initialization while recovering the teacherโs full-depth capability. We instantiate this setting on BoxPick-Exact-4 with GC-DQN. Since the scratch policy does not inherit the pretrained weights, we make two modifications to FlexLoop. First, at the full depth, we replace the standard TD target with the value estimate queried from the pretrained looped policy: L ํ TD =ํผ D ํ ํ ํพ (ํ , ํ, ํ)โ sg[ํฆ ํ ] 2 , ํฆ ํ = ํ+ ํพ(1โ ํ) max ํ โฒ ํ ํ (ํ โฒ , ํ, ํ โฒ ).(18) Second, we introduce an inter-depth Bellman refinement objective for shallow recurrent steps. Using the FlexLoopโs target network ฬ ํ, we define ํฆ ํ = ํ+ ํพ(1โ ํ) max ํ โฒ ํ ฬ ํ ํ (ํ โฒ , ํ, ํ โฒ ),(19) 24 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop-5 (Scratch) BoxPick-Exact-4 k=1k=2k=3k=4k=5 (a) Training Tasks. 0.00.10.20.30.40.5 Env Steps (B) 0.0 0.2 0.4 0.6 Success Rate IRU-5 Pretraining 0.00.10.20.30.40.5 Env Steps (B) FlexLoop-5 (Scratch) BoxPick-Exact-4 k=1k=2k=3k=4k=5 (b) Unseen Tasks. 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. and optimize L DQN aux = 1 ํพโ 2 ํพโ2 โ๏ธ ํ=1 ํผ D h ํ ํ ํ ํ+1 (ํ , ํ, ํ)โ sg[ํฆ ํ ] i .(20) We set the weight ofL DQN aux to 0.1. Notably, these modifications are introduced only to compensate for the absence of pretrained initial- ization. We observe no additional benefit in the teacher-initialized setting. As shown in Figures 13 and 14, FlexLoop from scratch rapidly develops depth elasticity and recovers strong full-depth performance, while still supporting effective state-wise adaptive inference. Together, these results broaden the applicability of FlexLoop to settings where only output-level access to a pretrained policy is available. D.2. FlexLoop on Policies with Different Pretraining Depths We further apply FlexLoop to IRU policies pretrained with different recurrent depths,ํพ โ 4,5,8,10. Figure 15 shows that, across all pretraining depths, FlexLoop substantially improves shallow-depth performance while preserving full-depth capability, with the same trend observed on unseen tasks. The corresponding Pareto curves further show that state-wise adaptive inference remains effective across different pretraining depths. We note that deeper pretrained policies generally require smaller stopping thresholdsํ, suggesting that the appropriate threshold scale depends on the pretrained recurrent depth. 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-4 FlexLoop-5 (Scratch)IRU-5 (a) Elasticity (Training). 12345 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-4 FlexLoop-5 (Scratch)IRU-5 (b) Elasticity (Unseen). 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate BoxPick-Exact-4 FlexLoop-5 (Adaptive)IRU-5 (Fixed K) (c) Pareto (Training). 12345 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate BoxPick-Exact-4 FlexLoop-5 (Adaptive)IRU-5 (Fixed K) (d) Pareto (Unseen). 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. 25 Depth-Elastic Looped Policies for Adaptive Test-Time Computation in Deep RL 1234 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate K max = 4 12345 Recurrence Depth K max = 5 2468 Recurrence Depth K max = 8 246810 Recurrence Depth K max = 10 FlexLoop (Ours) IRU (Fixed K) (a) Elasticity on Training Tasks. 1234 Recurrence Depth 0.0 0.2 0.4 0.6 0.8 Success Rate K max = 4 12345 Recurrence Depth K max = 5 2468 Recurrence Depth K max = 8 246810 Recurrence Depth K max = 10 FlexLoop (Ours) IRU (Fixed K) (b) Elasticity on Unseen Tasks. 1234 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 1.0 Success Rate K max = 4 12345 Average Recurrent Depth K max = 5 2468 Average Recurrent Depth K max = 8 246810 Average Recurrent Depth K max = 10 FlexLoop (Adaptive) IRU (Fixed K) (c) Pareto on Training Tasks. 1234 Average Recurrent Depth 0.0 0.2 0.4 0.6 0.8 Success Rate K max = 4 12345 Average Recurrent Depth K max = 5 2468 Average Recurrent Depth K max = 8 246810 Average Recurrent Depth K max = 10 FlexLoop (Adaptive) IRU (Fixed K) (d) Pareto on Unseen Tasks. 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. 26