Offline primitive discovery for accelerating data-driven reinforcement learning

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.