IP Library Granted Patent US 12,675,670
Granted Patent B2
US 12,675,670 · App. 18/044,852 · Granted Jul 7, 2026

Offline primitive discovery for accelerating data-driven reinforcement learning

Inventors: Anurag Ajay (Cambridge, MA); Ofir Nachum (San Francisco, CA); Aviral Kumar (Berkeley, CA); Sergey Levine (Berkeley, CA)
Assignee: Google LLC
G06N3/0455G06N3/092
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,675,670
App. No.
18/044,852
Filed
Mar 10, 2023
Granted
Jul 7, 2026
Kind
B2
Art Unit
2129
USPC
706/15
Abstract

A method includes determining a first state associated with a particular task, and determining, by a task policy model, a latent space representation of the first state. The task policy model may have been trained to define, for each respective state of a plurality of possible states associated with the particular task, a corresponding latent space representation of the respective state. The method also includes determining, by a primitive policy model and based on the first state and the latent space representation of the first state, an action to take as part of the particular task. The primitive policy model may have been trained to define a space of primitive policies for the plurality of possible states associated with the particular task and a plurality of possible latent space representations. The method further includes executing the action to reach a second state associated with the particular task.

Claims (78)

1 . A computer-implemented method comprising:

determining a first state associated with a particular task;

determining, by a task policy model, a latent space representation of the first state, wherein the task policy model has been trained to define, for each respective state of a plurality of possible states associated with the particular task, a corresponding latent space representation of the respective state;

determining, by a primitive policy model and based on (i) the first state and (ii) the latent space representation of the first state, an action to take as part of the particular task, wherein the primitive policy model has been trained to define a space of primitive policies for the plurality of possible states associated with the particular task and a plurality of possible latent space representations, and wherein the primitive policy model has been trained by a process comprising:

determining an undirected experience dataset comprising one or more sub-trajectories each representing a corresponding transition through a corresponding plurality of states of a training task;

determining, by an encoder model, for each respective sub-trajectory of the one or more sub-trajectories, one or more parameters of a corresponding latent space distribution; and

training the primitive policy model based on the one or more parameters of the corresponding latent space distribution; and

executing the action to reach a second state associated with the particular task, wherein the task policy model has been trained as part of an offline reinforcement learning process that comprises:

determining a reward-labeled version of the undirected experience dataset comprising one or more reward-labeled sub-trajectories each representing (i) a corresponding transition through a corresponding plurality of states of the training task and (ii) one or more rewards associated with the corresponding transition;

determining, by the encoder model and for each respective reward-labeled sub-trajectory of the one or more reward-labeled sub-trajectories, a corresponding latent space representation; and

training the task policy model based on (i) the one or more reward-labeled sub-trajectories and (ii) the corresponding latent space representation of each respective reward-labeled sub-trajectory.

2 . The computer-implemented method of claim 1 , further comprising:

determining, by the primitive policy model and based on (i) the second state and (ii) the latent space representation of the first state, a second action to take as part of the particular task; and

executing the second action to reach a third state associated with the particular task.

3 . The computer-implemented method of claim 1 , further comprising:

determining a state trajectory comprising a plurality of states connected by a corresponding plurality of actions, wherein the plurality of states comprises c states, wherein the corresponding plurality of actions comprise c actions, and wherein determining the state trajectory comprises:

determining, by the primitive policy model, for each respective state of the c states, and based on the respective state and the latent space representation of the first state, a subsequent action of the c actions to take as part of the particular task, wherein the first state and the second state form part of the c states; and

executing, for each respective state of the c states, the subsequent action to reach a subsequent state of the c states.

4 . The computer-implemented method of claim 3 , wherein determining the state trajectory comprises:

determining a plurality of state trajectories each comprising a corresponding plurality of states connected by a corresponding plurality of actions, wherein the task policy model is configured to determine a respective latent space representation of a respective first state of each respective state trajectory of the plurality of state trajectories, and wherein the primitive policy model is configured to use the respective latent space representation to determine each of the corresponding plurality of actions of the respective state trajectory.

5 . The computer-implemented method of claim 1 , wherein the first state, the action, and the second state form part of a Markov decision process (MDP) representing the particular task.

6 . The computer-implemented method of claim 1 , wherein executing the action comprises one or more of (i) simulating the action or (ii) causing an agent to physically carry out the action.

7 . The computer-implemented method of claim 1 , wherein the space of primitive policies is continuous and temporally-extended.

8 . The computer-implemented method of claim 1 , wherein training of the primitive policy model based on the one or more parameters of the corresponding latent space distribution comprises:

determining, by the primitive policy model, based on the one or more parameters of the corresponding latent space distribution, and for each respective state of the corresponding plurality of states of the respective sub-trajectory, one or more parameters of a corresponding distribution of actions associated with the respective state;

determining a loss value using a loss function based on the one or more parameters of the corresponding distribution of the actions associated with each respective state of the corresponding plurality of states of the respective sub-trajectory; and

updating the primitive policy model based on the loss value.

9 . The computer-implemented method of claim 8 , wherein the corresponding transition of each respective sub-trajectory of the one or more sub-trajectories comprises a corresponding plurality of actions connecting the corresponding plurality of states, wherein the encoder model is configured to determine the one or more parameters of the corresponding latent space distribution based on the corresponding plurality of actions and the corresponding plurality of states, and wherein the primitive policy model is configured to determine the one or more parameters of the corresponding distribution of actions associated with the respective state further based on the respective state.

10 . The computer-implemented method of claim 8 , wherein the process by which the primitive policy model has been trained further comprises:

generating, by a prior model and for each respective sub-trajectory of the one or more sub-trajectories, second one or more parameters of the corresponding latent space distribution based on an initial state associated with the respective sub-trajectory; and

determining the loss value using the loss function further based on a comparison of (i) the one or more parameters of the corresponding latent space distribution determined by the encoder model and (ii) the second one or more parameters of the corresponding latent space distribution determined by the prior model.

11 . The computer-implemented method of claim 10 , wherein the loss function comprises an auto-encoding loss function configured to evaluate the one or more parameters of the corresponding distribution of the actions associated with the respective state and a Kullback-Leibler constraint configured to perform the comparison of (i) the one or more parameters of the corresponding latent space distribution determined by the encoder model and (ii) the second one or more parameters of the corresponding latent space distribution determined by the prior model, and wherein the loss function comprises a weighted sum of the auto-encoding loss function and the Kullback-Leibler constraint.

12 . The computer-implemented method of claim 1 , wherein the space of primitive policies represents actions that, based on the undirected experience dataset, are possible to execute as part of the particular task.

13 . The computer-implemented method of claim 1 , wherein the task policy model has been further trained using a Conservative Q-Learning process, wherein the primitive policy model has been updated using a behavioral cloning process based on (i) the one or more reward-labeled sub-trajectories and (ii) the corresponding latent space representation of each respective reward-labeled sub-trajectory, and wherein the primitive policy model is has been updated independently of the one or more rewards associated with the corresponding transition represented by each of the one or more reward-labeled sub-trajectories.

14 . The computer-implemented method of claim 1 , wherein the task policy model has been trained as part of a few-shot imitation learning process that comprises:

determining an expert-demonstrated experience dataset comprising one or more expert-demonstrated sub-trajectories each representing a corresponding transition through a corresponding plurality of states of the training task;

determining, by the encoder model and for each respective expert-demonstrated sub-trajectory of the one or more expert-demonstrated sub-trajectories, a corresponding latent space representation; and

training the task policy model based on (i) the one or more expert-demonstrated sub-trajectories and (ii) the corresponding latent space representation of each respective expert-demonstrated sub-trajectory.

15 . The computer-implemented method of claim 14 , wherein the task policy model has been trained using a behavioral cloning process, and wherein the primitive policy model has been updated using the behavioral cloning process based on (i) the one or more expert-demonstrated sub-trajectories and (ii) the corresponding latent space representation of each respective expert-demonstrated sub-trajectory.

16 . The computer-implemented method of claim 10 , wherein the task policy model has been trained as part of an online reinforcement learning process comprising:

selecting, by the task policy model, a particular latent space representation to be used for a predetermined number of time steps;

determining, by the primitive policy model based on the particular latent space representation, a plurality of training actions for the predetermined number of time steps;

executing the plurality of training actions to transition through a corresponding plurality of states; and

training the task policy model using a second loss function based on rewards associated with the corresponding plurality of states.

17 . The computer-implemented method of claim 16 , wherein the task policy model has been further trained by:

training the task policy model using a Soft Actor Critic (SAC) process;

generating, by the task policy model, third one or more parameters of the particular latent space representation;

generating, by the prior model and based on an initial state associated with the predetermined number of time steps, fourth one or more parameters of the particular latent space representation; and

training the task policy model further based on a comparison of (i) the third one or more parameters and (ii) the fourth one or more parameters, wherein the comparison of (i) the third one or more parameters and (ii) the fourth one or more parameters is based on a Kullback-Leibler constraint.

18 . A system comprising:

a processor; and

a non-transitory computer-readable medium having stored thereon instructions that, when executed by the processor, cause the processor to perform operations comprising:

determining a first state associated with a particular task;

determining, by a task policy model, a latent space representation of the first state,

wherein the task policy model has been trained to define, for each respective state of a plurality of possible states associated with the particular task, a corresponding latent space representation of the respective state;

determining, by a primitive policy model and based on (i) the first state and (ii) the latent space representation of the first state, an action to take as part of the particular task, wherein the primitive policy model has been trained to define a space of primitive policies for the plurality of possible states associated with the particular task and a plurality of possible latent space representations, and wherein the primitive policy model has been trained by a process comprising:

determining an undirected experience dataset comprising one or more sub-trajectories each representing a corresponding transition through a corresponding plurality of states of a training task;

determining, by an encoder model, for each respective sub-trajectory of the one or more sub-trajectories, one or more parameters of a corresponding latent space distribution; and

training the primitive policy model based on the one or more parameters of the corresponding latent space distribution; and

executing the action to reach a second state associated with the particular task, wherein the task policy model has been trained as part of an offline reinforcement learning process that comprises:

determining a reward-labeled version of the undirected experience dataset comprising one or more reward-labeled sub-trajectories each representing (i) a corresponding transition through a corresponding plurality of states of the training task and (ii) one or more rewards associated with the corresponding transition;

determining, by the encoder model and for each respective reward-labeled sub-trajectory of the one or more reward-labeled sub-trajectories, a corresponding latent space representation; and

training the task policy model based on (i) the one or more reward-labeled sub-trajectories and (ii) the corresponding latent space representation of each respective reward-labeled sub-trajectory.

19 . The system of claim 18 , wherein the operations further comprise:

determining a state trajectory comprising a plurality of states connected by a corresponding plurality of actions, wherein the plurality of states comprises c states, wherein the corresponding plurality of actions comprise c actions, and wherein determining the state trajectory comprises:

determining, by the primitive policy model, for each respective state of the c states, and based on the respective state and the latent space representation of the first state, a subsequent action of the c actions to take as part of the particular task, wherein the first state and the second state form part of the c states; and

executing, for each respective state of the c states, the subsequent action to reach a subsequent state of the c states.

20 . A non-transitory computer-readable medium having stored thereon instructions that, when executed by a computing device, cause the computing device to perform operations comprising:

determining a first state associated with a particular task;

determining, by a task policy model, a latent space representation of the first state, wherein the task policy model has been trained to define, for each respective state of a plurality of possible states associated with the particular task, a corresponding latent space representation of the respective state;

determining, by a primitive policy model and based on (i) the first state and (ii) the latent space representation of the first state, an action to take as part of the particular task, wherein the primitive policy model has been trained to define a space of primitive policies for the plurality of possible states associated with the particular task and a plurality of possible latent space representations, and wherein the primitive policy model has been trained by a process comprising:

determining an undirected experience dataset comprising one or more sub-trajectories each representing a corresponding transition through a corresponding plurality of states of a training task;

determining, by an encoder model, for each respective sub-trajectory of the one or more sub-trajectories, one or more parameters of a corresponding latent space distribution;

training the primitive policy model based on the one or more parameters of the corresponding latent space distribution; and

executing the action to reach a second state associated with the particular task, wherein the task policy model has been trained as part of an offline reinforcement learning process that comprises:

determining a reward-labeled version of the undirected experience dataset comprising one or more reward-labeled sub-trajectories each representing (i) a corresponding transition through a corresponding plurality of states of the training task and (ii) one or more rewards associated with the corresponding transition;

determining, by the encoder model and for each respective reward-labeled sub-trajectory of the one or more reward-labeled sub-trajectories, a corresponding latent space representation; and

training the task policy model based on (i) the one or more reward-labeled sub-trajectories and (ii) the corresponding latent space representation of each respective reward-labeled sub-trajectory.

Assignments (1)
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded Mar 13, 2023
From: AJAY, ANURAG; NACHUM, OFIR; KUMAR, AVIRAL; LEVINE, SERGEY
To: GOOGLE LLC
Reel/Frame 062960/0830 →
Continuity (2)
Provisional Application 63086986 · Oct 2, 2020
Related Publication 20230367996A1 · Nov 16, 2023
References Cited (49)
Shankar, Tanmay, and Abhinav Gupta. “Learning robot skills with temporal variational inference.” International Conference on Machine Learning. PMLR, 2020. (Year: 2020). [cited by examiner]
Lynch, Corey, et al. “Learning latent plans from play.” Conference on robot learning. Pmlr, 2020. (Year: 2020). [cited by examiner]
Achiam et al., “Constrained policy optimization,” In Proceedings of the 34th International Conference on Machine Learning, JMLR. Org, 2017, vol. 70, pp. 22-31 (arXiv:1705.10528v1 May 30, 2017). [cited by applicant]
Atkeson et al., “No Falls, No Resets: Reliable Humanoid Behavior in the DARPA Robotics Challenge,” In 2015 IEEE-RAS 15th International Conference on Humanoid Robots (Humanoids), IEEE, 2015, pp. 623-630. [cited by applicant]
Bacon et al., “The Option-Critic Architecture,” arXiv:1609.05140v2, In Thirty-First AAAI Conference on Artificial Intelligence, 2017, 9 pages. [cited by applicant]
Brown et al., Superhuman AI for Multiplayer Poker, Science, 2019, pp. 885-890, v. 365, No. 6456. [cited by applicant]
Co-Reyes et al., “Self-Consistent Trajectory Autoencoder: Hierarchical Reinforcement Learning with Trajectory Embeddings,” arXiv:1806.02813v1, Jun. 7, 2018, 11 pages. [cited by applicant]
Dasari et al., “RoboNet: Large-Scale Multi-Robot Learning,” arXiv:1910.11215v2, Jan. 2, 2020, 15 pages. [cited by applicant]
Dulac-Arnold et al., “Challenges of Real-World Reinforcement Learning,” arXiv:1904.12901v1, Apr. 29, 2019, 13 pages. [cited by applicant]
Eysenbach et al., “Diversity is All You Need: Learning Skills without a Reward Function,” arXiv:1802.06070V6, Oct. 9, 2018, 22 pages. [cited by applicant]
Fu et al., “D4RL: Datasets for Deep Data-Driven Reinforcement Learning,” arXiv:2004.07219v4, Feb. 6, 2021, 19 pages. [cited by applicant]
Fujimoto et al., “Off-Policy Deep Reinforcement Learning without Exploration,” arXiv:1812.02900v3, Aug. 10, 2019, 23 pages. [cited by applicant]
Haarnoja et al., “Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor,” arXiv:1801.01290v2, Aug. 8, 2018, 14 pages. [cited by applicant]
Hausman et al., “Learning an Embedding Space for Transferable Robot Skills,” In International Conference on Learning Representations (ICLR), 2018, 16 pages. [cited by applicant]
Higgins et al., “Beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework,” In International Conference on Learning Representations (ICLR), Nov. 4, 2016, 13 pages. [cited by applicant]
Jabri et al., “Unsupervised Curricula for Visual Meta-Reinforcement Learning,” In Advances in Neural Information Processing Systems, 2019, pp. 10519-10531 (arXiv:1912.04226v1). [cited by applicant]
Jaques et al., “Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog,” arXiv:1907.00456v2, Jul. 8, 2019, 16 pages. [cited by applicant]
Kamyar et al., “EMaQ: Expected-Max Q-Learning Operator for Simple Yet Effective Offline and Online RL,” arXiv:2007.11091v2, Jan. 13, 2021, 28 pages. [cited by applicant]
Kingma et al., “Adam: A Method for Stochastic Optimization,” arXiv:1412.6980v9, Jan. 30, 2017, 15 pages. [cited by applicant]
Krishnan et al., “DDCO: Discovery of Deep Continuous Options for Robot Learning from Demonstrations,” arXiv:1710.05421v2, Oct. 31, 2017, 21 pages. [cited by applicant]
Kumar et al., Learning Navigation Subroutines from Egocentric Videos, In Conference on Robot Learning, PMLR, 2020, pp. 617-626 (arXiv:1905.12612v2). [cited by applicant]
Kumar et al., “Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction,” In Neural Information Processing Systems (NeurIPS), arXiv:1906.00949v2, Nov. 25, 2019, 19 pages. [cited by applicant]
Kumar et al., “Conservative Q-Learning for Offline Reinforcement Learning,” arXiv:2006.04779v3, Aug. 19, 2020, 31 pages. [cited by applicant]
Levine et al., “End-to-End Training of Deep Visuomotor Policies,” Journal of Machine Learning Research, 2016, pp. 1334-1373, vol. 17, No. 1. [cited by applicant]
Levine et al., “Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems,” arXiv:2005.01643v3, Nov. 1, 2020, 43 pages. [cited by applicant]
Lynch et al., “Learning Latent Plans from Play,” In Conference on Robot Learning, 2020, pp. 1113-1132 (arXiv:1903.01973v2). [cited by applicant]
Matsushima et al., “Deployment-Efficient Reinforcement Learning via Model-Based Offline Optimization,” arXiv:2006.03647v2, Jun. 23, 2020, 21 pages. [cited by applicant]
Merel et al., “Neural Probabilistic Motor Primitives for Humanoid Control,” arXiv:1811.11711v2, (Conference Paper ICLR 2019) Jan. 15, 2019, 14 pages. [cited by applicant]
Nachum et al., “Near-Optimal Representation Learning for Hierarchical Reinforcement Learning,” arXiv:1810.01257v2, (Conference Paper ICLR 2019) Jan. 9, 2019, 18 pages. [cited by applicant]
Nachum et al., “Data-Efficient Hierarchical Reinforcement Learning,” In Advances in Neural Information Processing Systems, (arXiv:1805.08296v4), 2018, pp. 3303-3313. [cited by applicant]
Nachum et al., Multi-Agent Manipulation via Locomotion using Hierarchical Sim2Real. arXiv:1908.05224v2, Oct. 7, 2019, 14 pages. [cited by applicant]
Nachum et al., “Why Does Hierarchy (Sometimes) Work so Well in Reinforcement Learning?” arXiv:1909.10618v2, Dec. 30, 2019. [cited by applicant]
Parmas et al., “PIPPS: Flexible Model-Based Policy Search Robust to the Curse of Chaos,” arXiv:1902.01240v1, Feb. 4, 2019, 10 pages. [cited by applicant]
Peng et al., “MCP: Learning Composable Hierarchical Control with Multiplicative Compositional Policies,” (arXiv:1905.09808v1), In Advances in Neural Information Processing Systems, 2019, pp. 3686-3697. [cited by applicant]
Peters et al., “Relative Entropy Policy Search,” Proceedings of the Twenty-Fourth AAAI Conference on Artificial Intelligence, 2010, pp. 1607-1612. [cited by applicant]
Pinto et al., “Supersizing Self-supervision: Learning to Grasp from 50k Tries and 700 Robot Hours,” (arXiv:1509.06825v1), In 2016 IEEE international conference on robotics and automation (ICRA), IEEE, 2016, pp. 3406-341… [cited by applicant]
Schulman et al., “Proximal Policy Optimization Algorithms,” arXiv:1707.06347v2, Aug. 28, 2017, 12 pages. [cited by applicant]
Shankar et al., “Learning Robot Skills with Temporal Variational Inference,” arXiv:2006.16232v1, Jun. 29, 2020, 13 pages (Published as a conference paper at ICLR 2021). [cited by applicant]
Sharma et al., “Dynamics-Aware Unsupervised Discovery of Skills,” arXiv:1907.01657V2, Feb. 14, 2019 (Published as a conference paper at ICLR 2020), 21 pages. [cited by applicant]
Silver et al., “Mastering the game of Go with deep neural networks and tree search,” Nature, 2016, pp. 484-489, vol. 529, No. 7587 (Abstract only). [cited by applicant]
Stolle et al., “Learning Options in Reinforcement Learning,” Lecture Notes in Computer Science, 2002, pp. 212-223, vol. 2371, doi: 10.1007/3-540-45622-8. [cited by applicant]
Tesauro, Gerald, “Temporal Difference Learning and TD-Gammon,” Communications of the ACM, Mar. 1995, 16 pages, vol. 38, No. 3. [cited by applicant]
Van Hasselt et al., “Deep Reinforcement Learning with Double Q-learning,” arXiv:1509.06461v3, Dec. 8, 2015, 13 pages. [cited by applicant]
Vezhnevets et al., “FeUdal Networks for Hierarchical Reinforcement Learning,” arXiv:1703.01161v2, Mar. 6, 2017, 12 pages. [cited by applicant]
Wang et al., “Robust Imitation of Diverse Behaviors,” arXiv:1707.02747v2, In Advances in Neural Information Processing Systems, 2017, pp. 5320-5329. [cited by applicant]
Wu et al., “Behavior Regularized Offline Reinforcement Learning,” arXiv:1911.11361v1, Nov. 26, 2019, 25 pages. [cited by applicant]
Yu et al., “Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning,” arXiv:1910.10897v2, In Conference on Robot Learning, 2020, pp. 1094-1100. [cited by applicant]
International Searching Authority, International Search Report and Written Opinion mailed Jan. 28, 2022, issued in connection with International Patent Application No. PCT/US2021/071562, filed Sep. 23, 2021, 22 pages. [cited by applicant]
Puterman, Martin L., Markov Decision Processes: Discrete Stochastic Dynamic Programming, Wiley, Aug. 2014 (Abstract Only). [cited by applicant]