IP Library › Granted Patent US 12,639,307
Granted Patent B2
US 12,639,307 · App. 18/867,222 · Granted May 26, 2026

Large-scale retrieval augmented reinforcement learning

Inventors: Peter Conway Humphreys (London, GB); Arthur Clement Guez (London, GB)
Assignee: GDM Holding LLC
G06F16/24542
View Patent ↗
Loading inventors, assignments & file history…
Monitor This Case
Get email alerts when status or documents change.
Order Certified Copies
Most orders are placed with the USPTO same day — all within 24 business hours.
Order via The Patent Place →
Pre-filled with this patent's details
Quick Facts
Patent No.
US 12,639,307
App. No.
18/867,222
Granted
May 26, 2026
Kind
B2
Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for controlling a reinforcement learning agent in an environment to perform a task. In one aspect, a method comprises: maintaining a retrieval dataset that stores a plurality of history observations and, for each history observation, a respective associated context; receiving a current observation characterizing a current state of the environment; selecting one or more history observations from the plurality of history observations; processing, using an encoder neural network and in accordance with current values of encoder network parameters, an encoder network input comprising (i) the current observation and (ii) the one or more selected history observations and their respective associated context to generate a latent state representation for the current state of the environment; and using the latent state representation to determine an action to be performed by the agent in response to the current observation.

Claims (64)

1 . A method for controlling an agent in an environment to perform a task, the method comprising:

maintaining a retrieval dataset that stores a plurality of history observations and, for each history observation, a respective associated context, wherein each history observation characterizes a respective history state of the environment and the respective associated context for each history observation characterizes one or more actions that were performed in response to the history observation;

receiving a current observation characterizing a current state of the environment;

selecting one or more history observations from the plurality of history observations, comprising:

generating a query vector for the current observation, comprising:

processing the current observation using an embedding neural network to generate an embedding and processing the embedding to generate the query vector, and

selecting the one or more history observations from the plurality of history observations using the query vector; and

processing, using an encoder neural network and in accordance with current values of encoder network parameters, an encoder network input comprising (i) the current observation and (ii) the one or more selected history observations and their respective associated context to generate a latent state representation for the current state of the environment;

processing, by a policy neural network, an input that comprises the latent state representation generated by the encoder neural network to generate a policy network output; and

determining an action to be performed by the agent in response to the current observation based on the policy network output.

2 . The method of claim 1 , wherein the associated context of each history observation comprises one or both of: a history action performed in response to the history observation, or a history reward received in response to the history action being performed.

3 . The method of claim 1 , wherein the associated context of each history observation further comprises one or both of: a history final observation characterizing a terminal state in a trajectory that includes the history observation, a history final reward received by the agent in the trajectory that includes the history observation.

4 . The method of claim 1 , wherein processing the embedding to generate the query vector comprises:

processing the embedding by using principal component analysis.

5 . The method of claim 1 , wherein each history observation stored in the retrieval dataset is associated with a corresponding key vector that is generated by using the embedding neural network.

6 . The method of claim 5 , wherein selecting the one or more history observations from the plurality of history observations comprises:

selecting, from among the key vectors associated with the plurality of history observations, one or more key vectors based on their respective distances from the query vector for the current observation; and

selecting, as the one or more selected history observations, corresponding history observations from the plurality of history observations that are associated with the one or more selected key vectors.

7 . The method of claim 6 , wherein selecting the one or more key vectors comprises:

selecting one or more key vectors having smallest Euclidean distances from the query vector for the current observation.

8 . The method of claim 1 , wherein selecting the one or more history observations from the plurality of history observations using the query vector comprises using an approximate nearest neighbor matching technique.

9 . The method of claim 1 , wherein selecting the one or more history observations from the plurality of history observations using the query vector comprises using a vector similarity search technique.

10 . The method of claim 9 , wherein the vector similarity search technique comprises a maximum inner product search technique.

11 . The method of claim 9 , wherein the maximum inner product search technique comprises an anisotropic quantization-based maximum inner product search technique.

12 . The method of claim 1 , wherein some or all of the plurality of history observations stored in the retrieval dataset are generated as a result of the agent performing one or more different tasks when interacting with the environment.

13 . The method of claim 1 , wherein some or all of the plurality of history observations stored in the retrieval dataset each have a different dimension, a different modality, or both than the current observation.

14 . The method of claim 1 , wherein determining the action to be performed by the agent in response to the current observation comprises:

performing a look-ahead search that begins from the current state and searches through a latent state representation space guided by the policy network output and additional policy network outputs of the policy neural network, wherein the policy neural network is configured to, at each of multiple time steps, process a respective latent state representation for the time step to generate a corresponding additional policy network output for the time step.

15 . The method of claim 14 , wherein the look ahead search comprises a Monte-Carlo Tree Search (MCTS) technique.

16 . The method of claim 1 , wherein the policy network output comprises one or both of:

a predicted policy output that defines a score distribution over a set of possible actions; or

a predicted value output that defines a value representing a predicted long-term reward received by the agent as a result of the environment being in the current state.

17 . The method of claim 1 , wherein determining the action to be performed by the agent in response to the current observation comprises:

selecting, from the set of possible actions, the action to be performed by the agent by using the policy network output of the policy neural network.

18 . The method of claim 1 , wherein the task comprises one of: an object manipulation task or an environment navigation task.

19 . The method of claim 1 , wherein the agent is a mechanical agent, the environment is a real-world environment, and the observation comprises data from one or more sensors configured to sense the real-world environment.

20 . The method of claim 1 , wherein the encoder neural network comprises:

a first sub network configured to process the current observation to generate a first encoded representation;

a second sub network configured to process each of the one or more selected history observations and its respective associated context to generate a respective second encoded representation; and

a third sub network configured to process the first encoded representation and the respective second encoded representations to generate the latent state representation.

21 . The method of claim 1 , wherein the encoder neural network has been trained through reinforcement learning on a training dataset that stores a plurality of trajectories generated as a result of the reinforcement learning agent interacting with the environment.

22 . The method of claim 21 , wherein the retrieval dataset is the same as the training dataset.

23 . The method of claim 21 , wherein the retrieval dataset is different from the training dataset and is larger than the training dataset.

24 . The method of claim 21 , wherein the encoder neural network has been trained through model-based reinforcement learning.

25 . One or more non-transitory computer storage media storing instructions that when executed by one or more computer cause the one or more computers to perform operations for controlling an agent in an environment to perform a task, the operations comprising:

maintaining a retrieval dataset that stores a plurality of history observations and, for each history observation, a respective associated context, wherein each history observation characterizes a respective history state of the environment and the respective associated context for each history observation characterizes one or more actions that were performed in response to the history observation;

receiving a current observation characterizing a current state of the environment;

selecting one or more history observations from the plurality of history observations, comprising:

generating a query vector for the current observation, comprising:

processing the current observation using an embedding neural network to generate an embedding and processing the embedding to generate the query vector, and

selecting the one or more history observations from the plurality of history observations using the query vector; and

processing, using an encoder neural network and in accordance with current values of encoder network parameters, an encoder network input comprising (i) the current observation and (ii) the one or more selected history observations and their respective associated context to generate a latent state representation for the current state of the environment;

processing, by a policy neural network, an input that comprises the latent state representation generated by the encoder neural network to generate a policy network output; and

determining an action to be performed by the agent in response to the current observation based on the policy network output.

26 . A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform operations for controlling an agent in an environment to perform a task, the operations comprising:

maintaining a retrieval dataset that stores a plurality of history observations and, for each history observation, a respective associated context, wherein each history observation characterizes a respective history state of the environment and the respective associated context for each history observation characterizes one or more actions that were performed in response to the history observation;

receiving a current observation characterizing a current state of the environment;

selecting one or more history observations from the plurality of history observations, comprising:

generating a query vector for the current observation, comprising:

processing the current observation using an embedding neural network to generate an embedding and processing the embedding to generate the query vector, and

selecting the one or more history observations from the plurality of history observations using the query vector; and

processing, using an encoder neural network and in accordance with current values of encoder network parameters, an encoder network input comprising (i) the current observation and (ii) the one or more selected history observations and their respective associated context to generate a latent state representation for the current state of the environment;

processing, by a policy neural network, an input that comprises the latent state representation generated by the encoder neural network to generate a policy network output; and

determining an action to be performed by the agent in response to the current observation based on the policy network output.

Assignments (3)
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded Jun 6, 2025
From: DEEPMIND TECHNOLOGIES LIMITED
To: GDM HOLDING LLC
Reel/Frame 071498/0210 →
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded May 27, 2025
From: HUMPHREYS, PETER CONWAY; GUEZ, ARTHUR CLEMENT
To: DEEPMIND TECHNOLOGIES LIMITED
Reel/Frame 071225/0686 →
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded May 13, 2025
From: HUMPHREYS, PETER CONWAY; GUEZ, ARTHUR CLEMENT
To: DEEPMIND TECHNOLOGIES LIMITED
Reel/Frame 071102/0353 →
Continuity (2)
Provisional Application 63344008 · May 19, 2022
Related Publication 20250335439A1 · Oct 30, 2025
References Cited (53)
US 20190354869A1 · Warde-Farley et al. · 2019 [cited by applicant]
US 20200125586A1 · Rezaeian · 2020 [cited by examiner]
US 20200195652A1 · Carnahan · 2020 [cited by examiner]
US 20200341976A1 · Aggarwal · 2020 [cited by examiner]
US 20210182697A1 · Singh · 2021 [cited by examiner]
US 20220277031A1 · Quamar · 2022 [cited by examiner]
US 20230118240A1 · Wong · 2023 [cited by examiner]
Aytar et al., “Playing hard exploration games by watching YouTube,” Paper, Presented at the Thirty-Second Annual Conference on Neural Information Processing Systems (NIPS), Montreal, Canada, Dec. 2-8, 2018; Advances in … [cited by applicant]
AssaAbloy.com [online], “The DeepMind JAX Ecosystem,” Oct. 1, 2023, retrieved on Jan. 15, 2025, retrieved from URL <https://github.com/google-deepmind>, 3 pages. [cited by applicant]
Atkeson et al., “Locally weighted learning for control,” Artificial Intelligence Review, 1997, 11:75-113. [cited by applicant]
Barreto et al., “Practical Kernel-Based Reinforcement Learning,” The Journal of Machine Learning Research, Apr. 2016, 17(1):2372-2441. [cited by applicant]
Baudiš et al., “PACHI: State of the Art Open Source Go Program,” Paper, Presented at the 13th International Conference, Tilburg, The Netherlands, Nov. 20-22, 2011; Revised Selected Papers Advances in Computer Games, 201… [cited by applicant]
Blundell et al., “Model-Free Episodic Control,” CoRR, Submitted on Jun. 14, 2016, arXiv:1606.04460v1, 12 pages. [cited by applicant]
Borgeaud et al., “Improving language models by retrieving from trillions of tokens,” CoRR, Submitted on Dec. 8, 2021, arXiv:2112.04426v1, 43 pages. [cited by applicant]
Coulom, “Efficient Selectivity and Backup Operators in Monte-Carlo Tree Search,” Paper, Presented at the 5th International Conference on Computers and Games, Turin, Italy, May 29-31, 2006; Computers and Games, 2007, pp.… [cited by applicant]
De Haan et al., “Causal Confusion in Imitation Learning,” Paper, Presented at the Thirty-Third Annual Conference on Neural Information Processing Systems (NIPS), Vancouver, Canada, Dec. 8-14, 2019; Advances in Neural In… [cited by applicant]
Goyal et al., “Retrieval-Augmented Reinforcement Learning,” CoRR, Mar. 9, 2022, arXiv:2202.08417v2, 27 pages. [cited by applicant]
Guo et al., “Accelerating Large-Scale Inference with Anisotropic Vector Quantization,” Paper, Presented at the Proceedings of the 37th International Conference on Machine Learning, Virtual, Jul. 13-18, 2020; PMLR, 2020,… [cited by applicant]
Guu et al., “REALM: Retrieval-Augmented Language Model Pre-Training,” CoRR, Submitted on Feb. 10, 2020, arXiv:2002.08909v1, 12 pages. [cited by applicant]
Hansen et al., “Fast deep reinforcement learning using online adjustments from the past,” Paper, Presented at the Thirty-Second Annual Conference on Neural Information Processing Systems, Montreal, Canada, Dec. 2-8, 201… [cited by applicant]
He et al., “Identity Mappings in Deep Residual Networks,” Paper, Presented at the 14th European Conference on Computer Vision, Amsterdam, The Netherlands, Oct. 11-14, 2016; Computer Vision—ECCV 2016, Springer, 2016, pp.… [cited by applicant]
International Preliminary Report on Patentability in International Appln. No. PCT/EP2023/063492, dated Nov. 7, 2024, 9 pages. [cited by applicant]
International Search Report and Written Opinion in International Appln. No. PCT/EP2023/063492, dated Jul. 24, 2023, 15 pages. [cited by applicant]
Jain et al., “Semiparametric Reinforcement Learning,” Workshop Submission, Presented at the Sixth International Conference on Learning Representations, Vancouver, Canada, Apr. 30-May 3, 2018; OpenReview, Feb. 2018, 6 pa… [cited by applicant]
Jouppi et al., “In-Datacenter Performance Analysis of a Tensor Processing Unit,” Paper, Presented at the Proceedings of the 44th Annual International Symposium on Computer Architecture, Toronto, Canada, Jun. 24-28, 2017… [cited by applicant]
Kigma et al., “Adam: A Method for Stochastic Optimization,” CoRR, Submitted on Dec. 22, 2014, arXiv:1412.6980v1, 9 pages. [cited by applicant]
Lampinen et al., “Towards mental time travel: a hierarchical memory for reinforcement learning agents,” Paper, Presented at the Thirty-Fifth Annual Conference on Neural Information Processing Systems, Virtual, Dec. 6-14… [cited by applicant]
Lengyel et al., “Hippocampal Contributions to Control: The Third Way,” Paper, Presented at the Twenty-First Annual Conference on Neural Information Processing Systems, Vancouver, Canada, Dec. 3-8, 2007; Advances in Neur… [cited by applicant]
Levine et al., “Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems,” CoRR, Submitted on Nov. 1, 2020, arXiv:2005.01643v3, 43 pages. [cited by applicant]
Lin et al., “Episodic Memory Deep Q-Networks,” CoRR, Submitted on May 19, 2018, arXiv:1805.07603v1, 7 pages. [cited by applicant]
Lin, “Self-improving reactive agents based on reinforcement learning, planning and teaching,” Machine Learning, May 1992, 8:293-321. [cited by applicant]
Martin et al., “The kNN-TD Reinforcement Learning Algorithm,” Paper, Presented at the IWINAC '09: Proceedings of the 3rd International Work-Conference on The Interplay Between Natural and Artificial Computation, Santiag… [cited by applicant]
Mnih et al., “Human-level control through deep reinforcement learning,” Nature, Feb. 2015, 518(7540):529-533. [cited by applicant]
Oh et al., “Control of Memory, Active Perception, and Action in Minecraft,” Paper, Presented at the Proceedings of The 33rd International Conference on Machine Learning, New York, New York, Jun. 19-24, 2016; PMLR, 2016,… [cited by applicant]
Ormoneit et al., “Kernel-Based Reinforcement Learning,” Machine Learning, Nov. 2002, 49:161-178. [cited by applicant]
Parisotto et al., “Stabilizing Transformers for Reinforcement Learning,” Paper, Presented at the Thirty-Seventh International Conference on Machine Learning, Virtual, Jul. 12-18, 2020; PMLR, 2020, pp. 7487-7498. [cited by applicant]
Pritzel et al., “Neural Episodic Control,” Paper, Presented at the Proceedings of the 34th International Conference on Machine Learning, Sydney, Australia, Aug. 6-11, 2017; PMLR, 2017, 70:2827-2836. [cited by applicant]
Ram et al., “Continuous case-based reasoning,” Artificial Intelligence, Feb. 1997, 90(1-2):25-77. [cited by applicant]
Reed et al., “A Generalist Agent,” CoRR, Submitted on Nov. 11, 2022, arXiv:2205.06175v3, 42 pages. [cited by applicant]
Riedmiller, “Neural Fitted Q Iteration—First Experiences with a Data Efficient Neural Reinforcement Learning Method,” Paper, Presented at the 16th European Conference on Machine Learning, Porto, Portugal, Oct. 3-7, 2005… [cited by applicant]
Ritter et al., “Been There, Done That: Meta-Learning with Episodic Recall,” Paper, Presented at the Proceedings of the 35th International Conference on Machine Learning, Stockholm, Sweden, Jul. 10-15, 2018; PMLR, 2018, … [cited by applicant]
Ritter et al., “Rapid Task-Solving in Novel Environments,” CoRR, Submitted on Jun. 5, 2020, arXiv:2006.03662v1, 16 pages. [cited by applicant]
Rosin, “Multi-armed bandits with episode context,” Annals of Mathematics and Artificial Intelligence, Aug. 2011, 61:203-230. [cited by applicant]
Schrittwieser et al., “Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model,” CoRR, Submitted on Feb. 21, 2020, arXiv:1911.08265v2. [cited by applicant]
Schrittwieser et al., “Mastering Atari, Go, chess and shogi by planning with a learned model,” Nature, Dec. 2020, 588(7839):604-609. [cited by applicant]
Schrittwieser et al., “Online and Offline Reinforcement Learning by Planning with a Learned Model,” Paper, Presented at the Thirty-Fifth Annual Conference on Neural Information Processing Systems, Virtual, Dec. 6-14, 20… [cited by applicant]
Schulman et al., “Trust Region Policy Optimization,” Paper, Presented at the Proceedings of the 32nd International Conference on Machine Learning, Lille, France, Jul. 6-11, 2015; PMLR, 2015, 37:1889-1897. [cited by applicant]
Shrestha et al., “DeepAveragers: Offline Reinforcement Learning by Solving Derived Non-Parametric MDPs,” CoRR, Submitted on Oct. 18, 2020, arXiv:2010.08891v1, 21 pages. [cited by applicant]
Silver et al., “A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play,” Science, Dec. 2018, 362(6419):1140-1144. [cited by applicant]
Sprechmann et al., “Memory-based Parameter Adaptation,” CoRR, Submitted on Feb. 28, 2018, arXiv:1802.10542v1, 16 pages. [cited by applicant]
Wayne et al., “Unsupervised Predictive Memory in a Goal-Directed Agent,” CoRR, Submitted on Mar. 28, 2018, arXiv:1803.10760v1, 57 pages. [cited by applicant]
Xu et al., “VLM: Task-agnostic Video-Language Model Pre-training for Video Understanding,” CoRR, Submitted on Sep. 30, 2021, arXiv:2105.09996v3, 13 pages. [cited by applicant]
Office Action in European Appln. No. 23726530.1, dated Mar. 12, 2026, 10 pages. [cited by applicant]