{"id":"2716510c-b681-4475-ad3d-c16a7f2d9a17","arxiv_id":"2506.15446","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Behavior foundation models with GRU memory outperform memory-free zero-shot RL baselines in most partially observable ExORL settings, but the advantage is inconsistent on Cheetah.","lead":"This paper adds memory modules such as GRUs to zero-shot reinforcement learning agents that were built for fully observable environments, and tests them when the agent only sees noisy or incomplete observations. On partial-observation versions of standard robotics benchmarks, the memory-augmented method usually beats memory-free baselines, though it is worse on some Cheetah tasks.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Equation 5's trajectory-level FB guarantee is assumed, not proven, and the paper's own Section 6.1 reports degeneracy in the combined F/B representation, so the central theoretical claim lacks support.","rationale":"The reader's weakest assumption is exactly the load-bearing concern: the paper ports the FB representation learning guarantee from Markov states to GRU hidden states without proof or diagnostic. This is central because the claimed advantage of FB-M over memory-free BFMs depends on the trajectory-level successor measure being accurately approximated by the dot product of memory-conditioned forward and backward embeddings. The internal admission in Section 6.1 that the combined F(f_F(tau))^T B(f_B(tau)) representation is degenerate for non-GRU models strengthens the concern, showing that the approximation is not universally valid and may be specific to the chosen architecture and hyperparameters. The mixed empirical results, particularly on Cheetah where FB-M underperforms memory-free FB, mean the empirical evidence does not independently rescue the theory. A targeted diagnostic of Eq. 5 and Eq. 6 would settle whether the theoretical grounding holds. Since this concern is addressable and the reader already recommended CONDITIONAL, I recommend keeping the verdict unchanged rather than moving to reject or accept.","tokens_in":18134,"tokens_out":3749,"duration_ms":36905,"concrete_test":"Train FB-M on Walker flickering with GRU memory, then run a diagnostic: for held-out start trajectories tau_0 and sampled future trajectories tau, compute the empirical discounted occupancy M_hat(tau_0, tau) by Monte Carlo rollouts of pi_z, and compare with the model's prediction F(f_F(tau_0),z)^T B(f_B(tau)) using rank correlation or MSE. Separately, on the reward-labelled dataset, regress R(s) on B(f_B(tau)) and report the R^2. A poor fit would directly invalidate Eq. 5 or Eq. 6, confirming the concern.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Equation 5 asserts that the successor measure over length-L trajectories is approximated by F(f_F(tau_0),z)^T B(f_B(tau)), and Equation 6 uses this B to infer task embeddings. The paper provides no proof that the FB guarantee of Touati & Ollivier transfers to GRU hidden states; the hidden states are learned compressions of the last L observations, and nothing ensures they are sufficient statistics or that the learned F and B satisfy the FB fixed-point equations in trajectory space. The paper's own Section 6.1 reports that when both F and B use non-GRU memory models, training collapses because the combined representation is degenerate, implying Eq. 5 is fragile rather than a general guarantee. Additionally, Eq. 6 assumes test rewards are linear in B(f_B(tau)), a new assumption not justified by the trajectory-level successor measure. If these fail, FB-M has no theoretical grounding, and its mixed empirical performance (e.g., worse than memory-free FB on Cheetah) is unexplained.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper extends behaviour foundation models (BFMs), specifically the forward-backward (FB) representation, to partially observable RL by inserting GRU memory models into the forward model, backward model, and policy, producing FB-M (and a USF variant in Appendix B). It identifies state misidentification and task misidentification as the two failure mechanisms, proposes trajectory-level successor-measure equations (Eq. 5) and a z-inference rule (Eq. 6), and evaluates FB-M against FB, HILP, FB-stack, and an MDP oracle on noisy/flickering ExORL (Walker, Cheetah, Quadruped) and on dynamics interpolation/extrapolation. The main empirical finding is that FB-M improves over memory-free baselines in several settings, especially Quadruped and dynamics extrapolation, while the GRU compares favorably to transformer/S4 memory models; the paper also reports instability when both F and B use non-GRU memories.","tokens_in":18326,"tokens_out":6054,"duration_ms":57858,"significance":"If the empirical findings hold, this is a useful and timely contribution: it is, to my knowledge, the first systematic study of zero-shot RL under partial observability, it repurposes a standard single-task remedy (memory models) in a sensible way, and it ships open-source code. The evaluation is generally careful: 5 seeds, IQM aggregation, bootstrap confidence intervals, per-task tables, and an MDP oracle are all present. However, the theoretical grounding of the trajectory-level FB equations is assumed rather than proven, the reported benefits are not uniform across domains (Cheetah is a clear counterexample to the abstract's blanket wording), and the checkpoint-selection protocol may inflate absolute performance. The paper's value is therefore mostly empirical, and it would be strengthened by a more cautious framing and by additional validation of the successor-measure approximation.","major_comments":[{"comment":"The paper asserts that the FB guarantee of Touati and Ollivier transfers from Markov states to GRU hidden states, but no proof or formal condition is given. This is load-bearing because Eq. (5) defines the method. Furthermore, Section 6.1 reports that when both F and B use non-GRU memory models, training collapses and states that 'the combined representation M(...) ≈ F(f_F(τ))^T B(f_B(τ+)) is degenerate.' That observation is in direct tension with the assumed validity of Eq. (5) for GRU memory models. Please either provide conditions under which Eq. (5) holds for memory-based embeddings, or explicitly reframe the contribution as an empirical adaptation and soften the theoretical framing. A concrete diagnostic would be to compare F(f_F(τ0),z)^T B(f_B(τ)) against Monte Carlo estimates of the trajectory-level successor measure on held-out trajectories.","section":"Section 4.3, Eq. (5), and Section 6.1"},{"comment":"The abstract states that the proposed methods 'show improved performance over memory-free baselines.' Table 3 contradicts this for Cheetah: in the flickering setting FB-M scores 173±51 vs 182±25 for memory-free FB, and in the noisy setting 150±59 vs 213±53. The paper's own Section 5.2 correctly avoids claiming an FB-M advantage over FB on Cheetah, but the abstract and conclusion do not. Please qualify the claim by domain or provide an aggregate statistical test (e.g., a paired bootstrap over tasks and seeds) that supports the blanket statement.","section":"Abstract and Table 3"},{"comment":"The protocol selects, for each method, the checkpoint at which the all-task IQM aggregated across seeds is maximized. This is a form of test-time model selection and can inflate absolute scores; it also makes the 'near-oracle' claims in Section 5.2 (e.g., Quadruped flickering, 673±19 vs oracle 729±6) hard to interpret. Since all methods use the same rule, relative rankings may be roughly preserved, but I ask for a sensitivity check with a fixed training budget (e.g., the final checkpoint) to confirm the main conclusions, and a statement of how much the selection rule affects reported numbers.","section":"Section 5.1, Evaluation protocol"},{"comment":"The task-embedding inference rule z_test ≈ E[R(s) B(f_B(τ))] is stated without justification. In standard FB, Eq. (3) is justified by the assumption that the reward lies in the span of the backward embeddings B(s). For trajectories, the reward depends on the terminal Markov state s_L, and it is not shown that R(s_L) is well approximated as a linear function of B(f_B(τ)) across the trajectory distribution. Please provide a derivation or at least an empirical check (e.g., compare z inferred via Eq. (6) with the oracle z used at training time, or verify that reward prediction error is small on held-out labelled trajectories).","section":"Equation (6)"}],"minor_comments":[{"comment":"There are several typos: 'partially observability' should be 'partial observability', 'MuJoCu' should be 'MuJoCo', 'comparision' should be 'comparison', and 'evaled' should be 'evaluated'.","section":"Abstract and Section 5.1"},{"comment":"The sentence 'Our proposed approach (FB-M) outperforms this baseline in all settings except Walker where it performs similarly' is ambiguous because the previous sentence mentions both FB-stack and memory-free FB; 'this baseline' should be explicitly identified as FB-stack.","section":"Section 5.2"},{"comment":"The text says 'We consider 3 environments (three locomotion and one goal-directed)' but Table 1 lists only Walker, Quadruped, and Cheetah, all classified as locomotion. Please correct the mismatch.","section":"Section A.1"},{"comment":"The caption of Figure 8 says 'standard FP architecture' and the text says 'a standard FP architecture'; these should be 'FB'.","section":"Section C.3"},{"comment":"The phrase 'a fixed hidden state size of 32^2 = 1024 dimensions' is an unusual notation; simply saying 1024 dimensions would be clearer.","section":"Section 6.1"}],"recommendation":"major_revision","confidential_remarks":"The empirical contribution is real and the evaluation is largely careful, but the abstract and conclusion overstate the results on Cheetah, and the theoretical framing needs to be aligned with what is actually shown. The checkpoint-selection issue is common in this literature but should be checked for sensitivity. Overall, the paper is publishable after revision."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nWorth a look if you work on behavior foundation models. This is the first paper to add recurrent memory to FB-style zero-shot RL under partial observability, and it contributes a clean failure-mode decomposition (state vs. task misidentification) plus modified ExORL POMDP benchmarks that others can reuse. The empirical core is honest: 5 seeds, IQM, bootstrap confidence intervals, per-task tables, an MDP oracle, and a limitations section that does not hide bad results. On Quadruped, FB-M is close to oracle; on dynamics extrapolation it often beats memory-free FB and sometimes the oracle. That part is real and useful.\n\nThe weak spot is the theory, and it is visible in the paper's own text. Equation 5 asserts that the trajectory-level successor measure is approximated by F(f_F(tau),z)^T B(f_B(tau)) with GRU hidden states, but no proof or diagnostic shows that the FB guarantee transfers from Markov states to these learned compressions. Equation 6 adds a further assumption that the test reward is linear in B(f_B(tau)); that is not justified by the original FB derivation. The stress-test note points to Section 6.1, where the authors themselves report training collapse when both F and B use non-GRU memory models because the combined representation is degenerate. That is direct evidence that Eq. 5 is fragile, not a general guarantee. It does not kill the empirical contribution, but the paper is best read as an architectural fix with experimental support rather than a theoretical extension.\n\nThe other soft spot is Cheetah. FB-M is below memory-free FB on both partially observed state settings there, so the abstract's \"improved performance over memory-free baselines\" overstates the aggregate. A careful reader will catch it; the authors should qualify the claim or explain the failure.\n\nNothing circular: evaluation is against external baselines and held-out tasks, and hyperparameters were tuned on Walker but not on the test tasks. Citations look appropriate. Code is linked through a project page without a commit hash; that is minor.\n\nOverall, this deserves a serious referee. It is a solid empirical contribution to a narrow but active subfield, with honest limitations and reproducible enough detail. The theory needs either a proof, a bounded assumption, or a clear statement that Eq. 5 is an approximation. I would engage with it.\n\nRecommendation: send to peer review, with the expectation that the framing around Eq. 5 and the Cheetah discrepancy get addressed.","headline":"Solid empirical first step on memory-augmented zero-shot RL under partial observability, but the trajectory-level FB theory is assumed rather than proven and the Cheetah results undercut the abstract's overclaim.","tokens_in":18836,"tokens_out":1994,"would_cite":true,"duration_ms":20001,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Memory models restore zero-shot RL when states are only partially observed.","keywords":["zero-shot reinforcement learning","partial observability","forward-backward representations","successor measures","memory models","POMDP","GRU","dynamics generalization"],"falsifier":"In a small partially observed gridworld where the true trajectory occupancy can be computed exactly, train FB-M and compare its predicted successor measure F(f_F(τ_0^L), z)^T B(f_B(τ^L)) against the Monte Carlo occupancy of the learned policy; if the two diverge substantially while returns stay high, the memory-based extension is helping for reasons other than the claimed mechanism.","tokens_in":17905,"feed_emoji":"🧠","tokens_out":6531,"duration_ms":56944,"temperature":0.7,"pith_summary":"The paper asks what happens to zero-shot reinforcement learning when the agent no longer sees the full Markov state, and shows that standard behavior foundation models degrade sharply through two failure routes: they misidentify the state when the forward model and policy receive partial observations, and they misidentify the task when the backward model receives partial observations. To fix this, the paper adds memory models—gated recurrent units (GRUs) that compress a length-L history of observations and actions—to each of the forward model, backward model, and policy, and trains the forward-backward objective over trajectories instead of states. If the paper is right, the memory-augmented version, FB-M, recovers most of the lost zero-shot performance: near-oracle behavior on Quadruped, better-than-baseline scores on most partially observed Walker tasks, and gains over memory-free methods when dynamics change at test time. This matters because partial observability is the rule in real-world deployment, not the exception.","feed_headline":"GRU memory restores zero-shot RL under partial observability","feed_subtitle":"Adding memory to zero-shot RL agents recovers lost performance when states are hidden or shift.","key_machinery":"The central object is the memory-augmented forward-backward representation, in which three gated recurrent units (GRUs) f_F, f_B, and f_π turn length-L observation-action histories into hidden states, and the successor measure over trajectories is factorized as F(f_F(τ_0^L), z)^T B(f_B(τ^L)), with the policy choosing actions by maximizing F(f_π(τ^L), z)^T z. This dot product is what lets the agent both predict long-run occupancy (through F) and infer the task (through B), and the paper's empirical claim is that this factorization survives the move from states to memory hidden states, with the GRU hidden state acting as a learned approximate Markov state.","core_discovery":"On the paper's own terms, the central claim is that the failure of zero-shot reinforcement learning under partial observability can be traced to two identifiable mechanisms—state misidentification, where the forward model and policy receive observations that do not reveal the underlying Markov state, and task misidentification, where the backward model cannot recover the task embedding from partial observations. The paper's contribution is to replace the state arguments of the forward model F, backward model B, and policy π with the hidden states of separate memory models f_F, f_B, and f_π that compress length-L trajectories of observations and actions, and to assert that the forward-backward training objective survives this replacement: the trajectory-level successor measure $M^{{π_z}}$(τ_0^L, X) is approximated by ∫_X F(f_F(τ_0^L), z)^T B(f_B(τ^L)) ρ(dτ^L), and the policy maximizes F(f_π(τ^L), z)^T z. By evaluating against memory-free FB, the stacked-observation FB-stack, and HILP on ExORL benchmarks amended with noisy states, flickering states, and changed dynamics, the paper reports that FB-M improves over all baselines in most settings, approaches oracle performance on Quadruped, and slightly exceeds the MDP-trained oracle in the extrapolation-dynamics case on Cheetah and Quadruped.","pith_inferences":["If the trajectory-level FB guarantee holds, then in environments where a bounded observation history is a sufficient statistic of the hidden state, a suitable context length L should recover full-observability performance asymptotically, suggesting a practical rule: set L to the mixing time of the observation process.","The training collapse when both F and B use non-GRU memory hints that the factorization M ≈ F(f_F)^T B(f_B) is especially fragile to poorly conditioned hidden-state representations; probing the condition number or spectral properties of the hidden-state embeddings during training could predict such collapses before they happen.","A natural testable extension is to vary the context lengths of the forward and backward memories independently for tasks whose rewards depend on shorter histories than their dynamics, which the paper notes but does not evaluate."],"forward_implications":["FB-M outperforms memory-free FB and the stacking baseline on almost all partially observed state settings, reaching close to oracle performance on Quadruped.","On dynamics changes at test time, FB-M matches or beats all baselines in every environment, including the extrapolation regime where it can slightly exceed the MDP-trained oracle.","The proposed memory-based formulation applies equally to USF-based behavior foundation models, not just FB.","GRU memory models outperform transformer and S4d memory models in this setting, and using non-GRU memory for both forward and backward models causes training collapse.","Stacking recent observations (FB-stack) is an inconsistent remedy: it helps on Walker and Quadruped partially observed states but hurts on Cheetah and in most dynamics-change settings."],"supporting_citations":[{"why":"Supplies the forward-backward representation theory and the guarantee that zero-shot RL rests on; this paper extends that theory by moving from states to memory-model hidden states.","marker":"Touati & Ollivier (2021)"},{"why":"Establishes the zero-shot RL problem and the evaluation protocol on ExORL that the memory-based methods are measured against.","marker":"Touati et al. (2023)"},{"why":"Provides the HILP baseline, a state-of-the-art zero-shot RL method that FB-M is compared against.","marker":"Park et al. (2024b)"},{"why":"Defines successor measures, the object that Equation 5 re-expresses at the trajectory level.","marker":"Blier et al. (2021)"},{"why":"Provides the noisy and flickering state environments and the memory-based methods that inspire the architecture.","marker":"Meng et al. (2021)"},{"why":"Shows recurrent policies and critics are strong baselines for POMDPs, motivating the separate memory models for F, B, and π.","marker":"Ni et al. (2021)"},{"why":"Supplies the dynamics-generalization protocol (interpolation and extrapolation) that the paper adapts to zero-shot RL.","marker":"Packer et al. (2018)"},{"why":"Provides the ExORL benchmark, RND datasets, and training protocol used for all experiments.","marker":"Yarats et al. (2022)"},{"why":"POPGym evidence that GRUs are the most performant memory model, used to justify the memory-model choice.","marker":"Morad et al. (2023)"},{"why":"Defines the GRU architecture used as the memory model in FB-M.","marker":"Cho (2014)"}],"fun_headline_variants":["Memory fixes zero-shot RL when states are hidden or shift","GRU memory restores zero-shot RL under partial observability","Zero-shot RL fails without memory on partially observed tasks","Hidden states, shifting dynamics? Memory saves zero-shot RL"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The paper assumes, without proof, that the forward-backward guarantee still holds when the successor measure is computed from GRU hidden states rather than true Markov states—that is, that the dot product of the memory-compressed forward and backward embeddings accurately approximates the trajectory-level occupancy, and that the Q function this dot product defines is a valid critic for the policy.","fun_headline_variants_meta":{"raw":{"variants":["Memory fixes zero-shot RL when states are hidden or shift","GRU memory restores zero-shot RL under partial observability","Zero-shot RL fails without memory on partially observed tasks","Hidden states, shifting dynamics? Memory saves zero-shot RL"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00032,"raw_usage":{"total_tokens":1804,"prompt_tokens":949,"completion_tokens":855,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":565,"completion_tokens_details":{"reasoning_tokens":788}},"tokens_in":565,"tokens_out":855,"duration_ms":8060,"temperature":1.0,"reasoning_tokens":788,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T19:34:09.614701+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"In a small partially observed gridworld where the true trajectory occupancy can be computed exactly, train FB-M and compare its predicted successor measure F(f_F(τ_0^L), z)^T B(f_B(τ^L)) against the Monte Carlo occupancy of the learned policy; if the two diverge substantially while returns stay high, the memory-based extension is helping for reasons other than the claimed mechanism.","supporting_citations":[{"cited_title":"Learning one representation to optimize all rewards","cited_arxiv_id":null,"evidence_quote":"Supplies the forward-backward representation theory and the guarantee that zero-shot RL rests on; this paper extends that theory by moving from states to memory-model hidden states."},{"cited_title":"Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations, 2023","cited_arxiv_id":null,"evidence_quote":"Establishes the zero-shot RL problem and the evaluation protocol on ExORL that the memory-based methods are measured against."},{"cited_title":"Memory-based deep reinforcement learning for pomdps","cited_arxiv_id":null,"evidence_quote":"Provides the noisy and flickering state environments and the memory-based methods that inspire the architecture."}],"review_version":2}