Paper deep dive
Memory-Based Advantage Shaping for LLM-Guided Reinforcement Learning
Narjes Nourzad, Carlee Joe-Wong
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/20/2026, 11:08:10 PM
Summary
The paper introduces a Memory-Based Advantage Shaping method for Reinforcement Learning (RL) that integrates Large Language Model (LLM) guidance into a memory graph. This graph encodes subgoals and trajectories from both LLM priors and successful agent rollouts. A utility function derived from this graph shapes the advantage function, providing additional guidance to the critic without altering the reward signal. This approach improves sample efficiency and early learning in sparse-reward environments while minimizing reliance on continuous LLM supervision.
Entities (8)
Relation Signals (9)
utility function → shapes → advantage function
confidence 95% · This utility shapes the advantage function, providing the critic with additional guidance without altering the reward.
Memory-Based Advantage Shaping → uses → Memory Graph
confidence 95% · Central to the method is a memory graph co-constructed from the agent’s rollouts together with offline priors and infrequent online outputs from an LLM.
Memory-Based Advantage Shaping → improves → Sample Efficiency
confidence 93% · Preliminary experiments in benchmark environments show improved sample efficiency and faster early learning compared to baseline RL methods
Memory Graph → generates → utility function
confidence 92% · From this graph, we derive a utility function that evaluates how closely the agent's trajectories align with prior successful strategies.
Memory Graph → encodes → subgoals
confidence 90% · The graph provides a compact representation of subgoals and task-relevant knowledge that evolves during training
Memory Graph → encodes → Trajectories
confidence 90% · The graph provides a compact representation of subgoals and task-relevant knowledge that evolves during training
Memory-Based Advantage Shaping → reduces → LLM reliance
confidence 90% · Our method relies primarily on offline input and only occasional online queries, avoiding dependence on continuous LLM supervision.
Memory-Based Advantage Shaping → matches → LLM4Teach
confidence 88% · Here, our method matches the state-of-the-art LLM4Teach in both mean return and success rate
Memory-Based Advantage Shaping → outperforms → PPO
confidence 85% · Preliminary experiments in benchmark environments show improved sample efficiency and faster early learning compared to baseline RL methods
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:In environments with sparse or delayed rewards, reinforcement learning (RL) incurs high sample complexity due to the large number of interactions needed for learning. This limitation has motivated the use of large language models (LLMs) for subgoal discovery and trajectory guidance. While LLMs can support exploration, frequent reliance on LLM calls raises concerns about scalability and reliability. We address these challenges by constructing a memory graph that encodes subgoals and trajectories from both LLM guidance and the agent's own successful rollouts. From this graph, we derive a utility function that evaluates how closely the agent's trajectories align with prior successful strategies. This utility shapes the advantage function, providing the critic with additional guidance without altering the reward. Our method relies primarily on offline input and only occasional online queries, avoiding dependence on continuous LLM supervision. Preliminary experiments in benchmark environments show improved sample efficiency and faster early learning compared to baseline RL methods, with final returns comparable to methods that require frequent LLM interaction.
Tags
Links
- Source: https://arxiv.org/abs/2602.17931v1
- Canonical: https://arxiv.org/abs/2602.17931v1
Trouble viewing inline? Open PDF directly →
Full Text
11,896 characters extracted from source content.
Expand or collapse full text
Memory-Based Advantage Shaping for LLM-Guided Reinforcement Learning Narjes Nourzad1, Carlee Joe-Wong2 This work was conducted during an internship at Carnegie Mellon University. Abstract In environments with sparse or delayed rewards, reinforcement learning (RL) incurs high sample complexity due to the large number of interactions needed for learning. This limitation has motivated the use of large language models (LLMs) for subgoal discovery and trajectory guidance. While LLMs can support exploration, frequent reliance on LLM calls raises concerns about scalability and reliability. We address these challenges by constructing a memory graph that encodes subgoals and trajectories from both LLM guidance and the agent’s own successful rollouts. From this graph, we derive a utility function that evaluates how closely the agent’s trajectories align with prior successful strategies. This utility shapes the advantage function, providing the critic with additional guidance without altering the reward. Our method relies primarily on offline input and only occasional online queries, avoiding dependence on continuous LLM supervision. Preliminary experiments in benchmark environments show improved sample efficiency and faster early learning compared to baseline RL methods, with final returns comparable to methods that require frequent LLM interaction. Introduction Reinforcement learning (RL) has achieved state-of-the-art results across diverse domains (Nourzad et al. 2024; Liu et al. 2024), yet its limitations become evident in settings that more closely resemble real-world conditions. Sparse or delayed rewards, partial observability, and long decision horizons often result in poor sample efficiency and unstable training dynamics. Large language models (LLMs) can help address these issues by providing high-level priors, suggesting subgoals, and structuring exploration (Schoepp et al. 2025). At the same time, LLMs introduce new challenges. Their outputs may be unreliable, queries can be expensive, and continuous supervision risks diluting the role of environment-driven feedback. Thus, it is challenging to integrate LLM guidance in a way that enhances RL without creating excessive dependence. The goal is to ensure external signals complement rather than distort optimization. In this work, we introduce an approach that incorporates external guidance into RL models through advantage shaping. Central to the method is a memory graph co-constructed from the agent’s rollouts together with offline priors and infrequent online outputs from an LLM. The graph provides a compact representation of subgoals and task-relevant knowledge that evolves during training, reducing dependence on real-time LLM access. A utility signal derived from the memory graph is used to softly shape advantage estimates, guiding policy updates and improving early exploration. The influence of this shaping term becomes asymptotically negligible once the agent’s policy surpasses the usefulness of LLM-derived guidance, ensuring that the convergence guarantees of PPO (proximal policy optimization) (Schulman et al. 2017) remain intact. Methodology Figure 1: Overview of the proposed method. Our method integrates external guidance through a memory graph, initialized with offline LLM priors (e.g., partial trajectory and subgoal decompositions) and augmented with agent rollouts and occasional online queries. Online queries are triggered adaptively when the agent fails to extract useful guidance from memory for several consecutive episodes. The LLM is constrained to the same partial observability as the agent and returns short plans that are either added to the graph or used to bias action preferences through soft logit injection. Formally, the memory graph is given by =((o,a)τj,ζj,r^j)j=1N∪κℓ=1L∪g⊳.G= \ ((o,a)_ _j, _j, r_j ) \_j=1^N\;∪\; \ _ \_ =1^L\;∪\;\ g_ \. Each trajectory node j consists of a partial observation oτjo_ _j, action aτja_ _j, goal term ζj∈gj,κℓgj _j∈\ g_j,\, _ g_j\ indicating a final goal (gj)( g_j) or an abstract subgoal (κℓgj)( _ g_j), and estimated reward r^j r_j for the corresponding action sequence. The second set of nodes κℓ=1L\ _ \_ =1^L represents subgoals inferred from environment descriptions, while the final term g⊳\ g_ \ denotes the agent’s target goal(s). These nodes are connected through goal–subgoal relationships provided directly by the LLM. The graph evolves during training, with new nodes added when novel behaviors are observed and rarely accessed nodes pruned as obsolete to stay keep the structure compact while ensuring consistent guidance with little LLM reliance. For each state-action pair (ot,at)(o_t,a_t), utility is computed by evaluating the agent’s observed behavior against the memory graph. It leverages the same rollouts used for advantage estimation under the current policy. Each pair in the trajectory τ=(ot,at)t=1Tτ=(o_t,a_t)_t=1^T is compared to a stored trajectory node (m) from the memory graph, aligned with the environment layout in the current rollout. Formally, Ut≐r^m⋅ρ(g⊳,ζm)⋅((ot,at),(ot’,at’)τm).U_t r_m·ρ( g_ , _m)·s ((o_t,a_t),(o_t’,a_t’)_ _m ). The similarity term (⋅,⋅)s(·,·) measures action agreement and how closely the agent’s trajectory follows the path implied by the stored segment (e.g., position or directional overlap). The factor ρ(⋅,⋅)ρ(·,·) serves to downweight matches that are behaviorally similar but pursue different subgoals, computed as the Jaccard similarity between the goals in the rollout and those encoded in memory. This ensures that utility reflects both behavioral similarity and semantic alignment with successful prior strategies. Overview of the method in Fig. 1. Algorithm 1 Shaped PPO actor (changes in purple) Collect =(st,at,rt)D=\(s_t,a_t,r_t)\ using πθ _θ Compute AtA_t and UtU_t from rollouts A~t=At+ξtUt [rgb]0.578125,0.45703125,0.84765625 A_t=A_t+ [rgb]0.578125,0.45703125,0.84765625 _tU_t for epoch=1epoch=1 to K, minibatch ℬ⊂kB _k do rt(θ)=πθ(at|st)/πθk(at|st)r_t(θ)= _θ(a_t|s_t)/ _ _k(a_t|s_t) ℒshaped(πθ)L^shaped( _θ) =[min(rt,1±εk)A~t]=E [ (r_t,1\!±\! _k) [rgb]0.578125,0.45703125,0.84765625 A_t ] θ←θ+αθ∇θ←θ+ _θ _θ ℒshaped(πθ)L^shaped( _θ) end for We incorporate memory-derived utility into the policy update by augmenting the standard advantage term. The advantage function quantifies how favorable an action is relative to the average action at state, reinforcing those with higher-than-expected returns and suppressing those that fall short. However, during early training the critic is poorly calibrated due to limited exploration, often producing near-uniform or noisy value estimates (Henderson et al. 2018). As a result, the computed advantages provide weak learning signals, even when the behavior is directed toward the task, leading to inefficient or unstable updates. This issue is more pronounced in sparse-reward or delayed-feedback settings, where AtA_t tends to be close to zero for most initial training timesteps. To address this, we define a shaped advantage A~t=At+ξtUtwhere 0<ξt≤1. A_t=A_t+ _tU_t 0< _t≤ 1. When critic feedback is weak, UtU_t supplies additional direction aligned with task objectives, compensating for flat or noisy gradients and accelerating learning. As training progresses and AtA_t becomes more reliable, the contribution of UtU_t naturally becomes negligible. Because the utility enters additively and does not modify the reward signal, the optimization dynamics of PPO are preserved, along with its convergence guarantees. More broadly, this mechanism is compatible with policy gradient methods relying on advantage estimation, providing a general way to incorporate structured guidance into RL. Figure 2: Mean return on FrozenLake-8x8 (top) and Doorkey (bottom). Our method has the highest return. Experimental Results Experimental Setup We evaluate our method in two environments: FrozenLake (Gymnasium) and Doorkey (MiniGrid suit), selected to probe early-stage exploration effects and long-horizon, partially observable planning, respectively. We compare against three baselines: PPO, trained tabula rasa from rewards, Hierarchical RL (HRL) (Bhambri et al. 2024), which uses LLM-derived option policies for temporal abstraction, and LLM4Teach (Zhou et al. 2023), a state-of-the-art method where a pre-trained LLM distills guidance into a student RL agent. Figure 2 reports mean return across environments. In FrozenLake, shaped advantages achieve faster convergence, but baselines eventually reach the same asymptotic return with sufficient training. In DoorKey, however, the shaped agent sustains a clear advantage, achieving both higher sample efficiency and stronger final performance. While HRL improves over PPO, it trails our method by nearly a factor of two. Table 1 reports test-time evaluation on unseen seeds for DoorKey. Here, our method matches the state-of-the-art LLM4Teach in both mean return and success rate, with no statistically significant difference at the 95%95\% confidence level. Unlike teacher-based approaches, however, these gains are achieved without continuous LLM supervision, relying only on limited offline and occasional online queries. Method Mean Return Success Rate LLM4teach 0.912±0.0750.912± 0.075 0.970±0.0040.970± 0.004 Our method 0.898±0.0930.898± 0.093 0.953±0.0430.953± 0.043 Table 1: Performance comparison on unseen seeds. These results indicate that memory-derived utility provides meaningful guidance in sparse and partially observable environments while minimizing the reliance on expensive LLM queries. References S. Bhambri, A. Bhattacharjee, D. Kalwar, L. Guan, H. Liu, and S. Kambhampati (2024) Extracting heuristics from large language models for reward shaping in reinforcement learning. arXiv preprint arXiv:2405.15194. Cited by: Experimental Setup. P. Henderson, R. Islam, P. Bachman, J. Pineau, D. Precup, and D. Meger (2018) Deep reinforcement learning that matters. In Proceedings of the AAAI conference on artificial intelligence, Vol. 32. Cited by: Methodology. J. Liu, C. Chang, S. Sun, and T. Yu (2024) Integrating planning and deep reinforcement learning via automatic induction of task substructures. In The Twelfth International Conference on Learning Representations, Cited by: Introduction. N. Nourzad, J. Coleman, Z. Zhao, B. Krishnamachari, G. Verma, and S. Segarra (2024) Actor-twin framework for task graph scheduling. In The Seventeenth Workshop on Adaptive and Learning Agents, Cited by: Introduction. S. Schoepp, M. Jafaripour, Y. Cao, T. Yang, F. Abdollahi, S. Golestan, Z. Sufiyan, O. R. Zaiane, and M. E. Taylor (2025) The evolving landscape of llm-and vlm-integrated reinforcement learning. arXiv preprint arXiv:2502.15214. Cited by: Introduction. J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017) Proximal policy optimization algorithms. In arXiv preprint arXiv:1707.06347, External Links: Link Cited by: Introduction. Z. Zhou, B. Hu, C. Zhao, P. Zhang, and B. Liu (2023) Large language model as a policy teacher for training reinforcement learning agents. arXiv preprint arXiv:2311.13373. Cited by: Experimental Setup.