MIRA: Memory-Integrated Reinforcement Learning Agent with Limited LLM Guidance
Authors: Narjes Nourzad, Carlee Joe-Wong
Organizations: University of Southern California, Electrical and Computer Engineering, Los Angeles, CA 90007, USA · Carnegie Mellon University, Electrical and Computer Engineering, Pittsburgh, PA 15213, USA
Reinforcement learning (RL) agents often suffer from high sample complexity in sparse or delayed reward settings due to limited prior structure. Large language models (LLMs) can provide subgoal decompositions, plausible trajectories, and abstract priors that facilitate early learning. However, heavy reliance on LLM supervision introduces scalability constraints and dependence on potentially unreliable signals. We propose MIRA (Memory-Integrated Reinforcement Learning Agent), which incorporates a structured, evolving memory graph to guide early training. The graph stores decision-relevant information, including trajectory segments and subgoal structures, and is constructed from both the agent's high-return experiences and LLM outputs. This design amortizes LLM queries into a persistent memory rather than requiring continuous real-time supervision. From this memory graph, we derive a utility signal that softly adjusts advantage estimation to influence policy updates without modifying the underlying reward function. As training progresses, the agent's policy gradually surpasses the initial LLM-derived priors, and the utility term decays, preserving standard convergence guarantees. We provide theoretical analysis showing that utility-based shaping improves early-stage learning in sparse-reward environments. Empirically, MIRA outperforms RL baselines and achieves returns comparable to approaches that rely on frequent LLM supervision, while requiring substantially fewer online LLM queries. Project webpage: https://narjesno.github.io/MIRA/
Figures & tables
Figure 1: Offline priors and online LLM suggestions are filtered by a screening unit before being incorporated into the memory graph as healthy grafts . MIRA agent acts under partial observations, interacting with the environment. A utility module evaluates trajectory rollouts against the evolving memory graph , producing a utility signal that shapes advantage estimation and policy updates.
Figure 2: MIRA’s evolving memory graph. Trajectory segments τj are grouped under subgoal nodes κℓ . Subgoals can be shared across multiple final goals, enabling reuse of common behaviors.
Figure 3: Evaluation environments. Top: RedBall (navigation to target), LavaCrossing (long-horizon navigation with irreversible hazards), DoorKey (sparse reward with key–goal dependency). Bottom: RedBlueDoor (sequence-sensitive toggling), Distracted DoorKey (distractor-rich variant with key-goal dependency).
Figure 4: Mean return on FrozenLake (left): Both MIRA variants improve early-stage learning relative to PPO, while PPO eventually attains a comparable asymptotic return. Evolution of shaping terms ηt , ξt , and ratio δt (right): δt decays during training, ensuring convergence as δt→0 .
Figure 5: Mean return (top) and success rate (bottom) across four MiniGrid and BabyAI tasks. MIRA consistently outperforms both baselines, achieving faster learning, higher asymptotic return, and greater success rates. These results are obtained with a small LLM budget, using fewer than ten offline prompts to build memory graphs plus infrequent online queries to guide exploration.
Figure 6: Mean return (left): LLM4Teach shows faster early gains, while MIRA steadily improves and matches its final return with fewer queries. LLM-RS benefits early from reward shaping but plateaus lower. Return per LLM query (right): Under two query budgets, MIRA achieves higher efficiency.
Figure 7: Ablation Studies . Query frequency (left): Agents share the same offline memory but vary in online budgets. More queries accelerate learning, with high-budget agents achieving optimal returns more quickly. Unreliable LLM (middle): With identical offline memory, screening is disabled and queries are swapped from GPT-o4-mini to Gemini Pro only in the late phase. Performance remains stable in the late phase, indicating reduced dependence on online guidance once policy have matured. LLM models (right): Agents differ only in the LLM used for memory. Performance differences reflect divergent reasoning styles: Gemma3 induces inefficient checking, Claude favors cautious exploration, while Gemini Pro and o4-mini enable faster learning.
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Frozen Lake (Gymnasium)
Figure 9: Prompt and response from GPT-4o-mini for the FrozenLake environment. The LLM receives a serialized map or a textual description and is asked to provide the best and second-best safe trajectories from start to goal, avoiding all holes under deterministic dynamics.
Figure 10: Prompt to Offline LLM for the custom MiniGrid variant Distracted DoorKey . The prompt describes the task setting, object roles, and challenges, and asks the LLM to confirm understanding before suggestions.
Figure 11: Reasoning traces produced by different LLMs in response to our custom environment prompt as part of “Offline LLM” prompting. After confirming they understood the environment, each model was asked: “If you were to give an RL agent useful trajectories to help solve this, what would you do?” For models that do not output internal reasoning (e.g., Claude), we issued a follow-up prompt requesting their thought process. We omit repeated environment restatements and show only the key parts where the model explains how it decided on the action sequence.
Figure 12: Reasoning traces produced by Gemini and ChatGPT under “Online LLM” prompting. The prompt emphasizes that sufficient exploration has already been performed and, from the partial observation, no key is visible. A (flawed but plausible) line of reasoning is that the agent must still be in the phase of searching for the key, so reducing the probability of toggle appears reasonable to prioritize movement actions for exploration.
Figure 13: FrozenLake robustness to prompt wording. MIRA achieves similar performance under the original and alternative prompts, showing stability to natural variations in task description.
Figure 14: Sensitivity to the screening threshold. Lenient thresholds graft more candidate nodes into memory early, producing broader shaping during exploration but a slower overall improvement rate. Stricter thresholds delay graph growth yet yield sharper mid-training gains once high-confidence nodes appear. All settings converge to a narrow performance band.
Figure 15: Return curves for different η values under fixed ξ settings. Markers indicate the first iteration where performance begins to diverge, signaling when At starts to meaningfully affect learning. Early on, the critic signal is weak, and A~t is driven mostly by the utility term. When ξ is large enough, shaping accelerates the critic’s contribution by up to 50 iterations and leads to around 2.5× higher return compared to the unshaped case.
Figure 16: Wall-clock runtimes across environments. Time required to reach a 0.5 return (left): PPO reaches 0.5 only on Frozen Lake , while both MIRA variants converge across tasks. Runtime for 2k training steps (right): Online MIRA incurs extra overhead from initial LLM queries, but this cost reduces wasted exploration and leads to faster convergence in terms of overall wall time.
Figure 17: Memory size growth over training. Early growth is followed by convergence once trajectories become consistent, with final sizes increasing from RedBall to DoorKey to Distracted DoorKey in line with task complexity.
Method
SR90Return↑
Final Return↑
Final Step↓
Baseline
0 ± 0.002
0.009 ± 0.001
10362
MIRA (offline)
0.233 ± 0.087
0.295 ± 0.123
10351
MIRA (mid)
0.284 ± 0.065
0.902 ± 0.012
10257
MIRA (large)
0.851 ± 0.060
0.91 ± 0.013
9961
Appendix
Table 1: Performance on DoorKey . SR90Return is the mean return when success rate first exceeds 90%; Final Return is the return at the end of training; Final Step is the total environment steps. MIRA variants outperform the baseline in both early and final return, with MIRA (large) achieving the highest values while converging fastest.
Method
DoorKey
LavaCrossing
RedBlueDoor
RedBall
Baseline RL
0.018 ± 0.016
0.012 ± 0.027
0.044 ± 0.042
0.329 ± 0.205
HRL
0.852 ± 0.017
0.798 ± 0.090
0.830 ± 0.021
0.939 ± 0.046
LLM4Teach
0.912 ± 0.075
0.884 ± 0.100
0.901 ± 0.082
0.946 ± 0.051
MIRA
0.898 ± 0.093
0.855 ± 0.132
0.911 ± 0.077
0.942 ± 0.054
Appendix
Table 2: Mean return on unseen seeds across MiniGrid environments. MIRA achieves high and stable success, comparable to LLM4Teach, despite requiring substantially fewer LLM queries.
Method
DoorKey
LavaCrossing
RedBlueDoor
RedBall
Baseline RL
0.023 ± 0.017
0.017 ± 0.020
0.036 ± 0.043
0.539 ± 0.064
HRL
0.897 ± 0.013
0.841 ± 0.085
0.892 ± 0.012
0.956 ± 0.025
LLM4Teach
0.970 ± 0.004
0.931 ± 0.011
0.956 ± 0.003
0.958 ± 0.021
MIRA
0.953 ± 0.043
0.913 ± 0.077
0.944 ± 0.020
0.956 ± 0.036
Appendix
Table 3: Success rate on unseen seeds across MiniGrid environments. MIRA achieves consistently high success rates, matching LLM4Teach while requiring fewer queries, and outperforming baseline and HRL methods.
Metric
LLM4Teach
MIRA
t
p
95% CI
DoorKey (MR)
0.912±0.075
0.898±0.093
0.203
0.8495
[–0.181, 0.209]
DoorKey (SR)
0.970±0.004
0.953±0.043
0.682
0.5647
[–0.0885, 0.1225]
LavaCrossing (MR)
0.884±0.100
0.855±0.132
0.303
0.7778
[–0.2443, 0.3023]
LavaCrossing (SR)
0.931±0.011
0.913±0.077
0.401
0.7260
[–0.1681, 0.2041]
RedBlueDoor (MR)
0.901±0.082
0.911±0.077
–0.154
0.8851
[–0.1906, 0.1706]
RedBlueDoor (SR)
0.956±0.003
0.944±0.020
1.028
0.4081
[–0.0362, 0.0602]
Appendix
Table 4: Welch’s t-test comparing LLM4Teach and MIRA (MR: Mean Return - SR: Success Rate). None of the differences are statistically significant at α=0.05 .
Property
Value
Observation Type
Discrete
Horizon
200
Reward Sparsity
Sparse
Action Space
4 (tabular)
Dynamics
Slippery, irreversible
Appendix
Table 5: FrozenLake environment details.
Parameter
Value
Learning rate
1×10−4
Batch size
512
Mini-batch size
64
Number of epochs
4
Entropy coefficient
0.01
Discount factor ( γ )
0.99
Appendix
Table 6: Hyperparameters of FrozenLake
Property
Value
Observation Type
RGB
Reward Sparsity
Sparse and delayed
Action Space
7 (tabular)
View Size
7
Horizon
300
Appendix
Table 7: MiniGrid suite details.
Environment
Dynamics
RedBall
Reversible
RedBlueDoor
Irreversible
LavaCrossing
Irreversible
DoorKey
Subgoal seq.
Distracted DoorKey
+Visual distractors
Appendix
Table 8: MiniGrid environments and their dynamics.
Table 26Table 27
Figure 18: Convolutional encoder architecture used to process the agent’s 56×56×3 RGB observation in MiniGrid environments. The input passes through a series of Conv2D layers, reducing spatial dimensions while increasing channel depth. The final activation is flattened and fed to both policy and value heads. This encoder captures spatial layout, object presence, and agent-centric context for decision-making.
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.
Narjes Nourzad, Carlee Joe-Wong
University of Southern California · Carnegie Mellon University
Reinforcement learning has greatly advanced the capabilities of large language models, but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. With all techniques combined, LoGRA reduces average training memory usage by up to 45.7% across reasoning tasks without compromising performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the Molt library.
Effective training-time guidance is central to multi-agent reinforcement learning (MARL), yet remains difficult in sparse-reward settings where weak supervision limits coordination and policy improvement, and existing methods often require substantial domain expertise or manual design effort. Large language models (LLMs) provide a promising alternative for flexible learning-signal design, yet existing LLM-based methods remain largely single-agent-oriented, one-shot, or weakly validated for the evolving training dynamics of cooperative MARL. To address these limitations, we propose LLM-ALSO, an iterative LLM-driven adaptive learning-signal optimization framework for MARL. Rather than directly deploying LLM-generated rewards, LLM-ALSO decomposes adaptation into iterative diagnosis, proposal, and validation: a Critic LLM diagnoses stage-specific learning and coordination failures from sparse-return metrics and compact behavior evidence, a Generator LLM proposes candidate reward-shaping configurations conditioned on the diagnosis, and branch-validation feedback refines candidates before they affect the main training trajectory. Through short-horizon validation and stage-aware adaptation, LLM-ALSO promotes only validated updates into training, reducing the risk of unreliable LLM-generated modifications. Experiments on sparse-reward cooperative MARL tasks show that LLM-ALSO improves sparse-evaluation performance and learning efficiency.
Xiaoguang Wu, Zhi Zheng, Hui Xiong
University of Science and Technology of China · Department of Computer Science and Engineering, The Hong Kong University of Science and Technology, Hong Kong SAR, China · 2Thrust of Artificial Intelligence, The Hong Kong University of Science and Technology (Guangzhou)