Batched reinforcement learning
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for batched reinforcement learning. For example, the batched reinforcement learning techniques can be used to determine a control policy for a robot in simulation and the control policy can then be used to control the robot in the real world. In one aspect, a method includes obtaining a plurality of current observations, each current observation characterizing a current state of a respective environment replica; processing the current observations in parallel using the action selection neural network in accordance with current values of the network parameters to generate an action batch; obtaining a transition tuple batch including a respective transition tuple for each of the environment replicas, the respective transition tuple for each environment replica including: (i) a subsequent observation and (ii) a reward; and training the action selection neural network on the batch of transition tuples.
1 . A method of training an action selection neural network to select actions to be performed by an agent interacting with an environment, wherein the action selection neural network has a plurality of network parameters and is configured to receive an input observation and to process the input observation in accordance with the network parameters to generate a network output that defines an action to be performed by the agent in response to the input observation, wherein the training is performed by processing a computation graph comprising nodes representing neural network operations and edges representing flow of input and output data of the neural network operations, wherein the computation graph includes an inference subgraph and a training subgraph, and wherein the method comprises, for each of a plurality of steps:
obtaining a plurality of current observations, each current observation characterizing a current state of a respective one of a plurality of environment replicas, wherein each environment replica is maintained by a separate process;
generating an observation batch as a single tensor that combines the plurality of current observations of the plurality of environment replicas maintained by the separate processes;
processing the current observations of the plurality of environment replicas in the single tensor in parallel using the action selection neural network in accordance with current values of the network parameters to generate an action batch that includes, for each environment replica, a respective action to be performed by the agent in response to the current observation characterizing the current state of the environment replica, wherein the processing comprises executing the inference subgraph to perform batched inference in parallel for the action selection neural network on the current observations in the single tensor to generate a respective network output for each current observation and select a respective action from the respective network output;
generating a transition tuple batch comprising a respective transition tuple for each of the environment replicas, wherein the transition tuple batch is generated by available processors of a plurality of first processors and using the separate processes in parallel, wherein the generating the transition tuple batch comprises constraining each respective sperate process to progress each respective environment replica with a single transition step using the respective action specified in the action batch, the respective transition tuple for each environment replica comprising:
(i) a subsequent observation characterizing a subsequent state that the environment replica transitioned into as a result of the agent performing the respective action in the action batch for the environment replica, and
(ii) a reward generated as a result of the environment replica transitioning into the subsequent state; and
training the action selection neural network by a plurality of second processors, on the transition tuple batch to update the current values of the network parameters using a reinforcement learning technique, the training comprising executing the training subgraph using the plurality of second processors, wherein the training subgraph takes as input the transition tuple batch and the current values of the network parameters and applies the reinforcement learning technique to the training tuples in the batch to generate updated values of the network parameters.
2 . The method of claim 1 , wherein the generating the transition tuple batch comprising the respective transition tuple for each of the environment replicas comprises:
providing the respective action in the action batch to the separate process that maintains the environment replica corresponding to the respective action to cause the environment replica to transition into the subsequent state in parallel; and
obtaining, from each of the separate processes, the subsequent observation and the reward for the environment replica maintained inside of the separate process.
3 . The method of claim 2 , wherein the generating the transition tuple batch further comprises:
after the subsequent observation and the reward have been obtained from all of the separate processes, generating the transition tuple batch from data obtained from the separate processes.
4 . The method of claim 1 , wherein the generating the transition tuple batch comprising the respective transition tuple for each of the environment replicas comprises:
issuing respective calls in parallel to each of the separate processes with the actions for the environment replicas;
waiting until a subsequent observation and a reward are obtained from each of the separate processes in response to the respective calls; and
after determining that the subsequent observation and the reward have been obtained from each of the separate processes, generating the transition tuple batch using the obtained subsequent observations and rewards.
5 . The method of claim 1 , wherein the reinforcement learning technique is a proximal policy optimization (PPO) algorithm.
6 . The method of claim 1 , wherein the training the action selection neural network is performed prior to the progressing of each respective environment replica is performed again in the next single transition step.
7 . The method of claim 1 , wherein the generating the transition tuple batch and training the action selection neural network are performed in parallel.
8 . 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 training an action selection neural network to select actions to be performed by an agent interacting with an environment, wherein the action selection neural network has a plurality of network parameters and is configured to receive an input observation and to process the input observation in accordance with the network parameters to generate a network output that defines an action to be performed by the agent in response to the input observation, wherein the training is performed by processing a computation graph comprising nodes representing neural network operations and edges representing flow of input and output data of the neural network operations, wherein the computation graph includes an inference subgraph and a training subgraph, and wherein the operations comprise, for each of a plurality of steps:
obtaining a plurality of current observations, each current observation characterizing a current state of a respective one of a plurality of environment replicas, wherein each environment replica is maintained by a separate process;
generating an observation batch as a single tensor that combines the plurality of current observations of the plurality of environment replicas maintained by the separate processes;
processing the current observations of the plurality of environment replicas in the single tensor in parallel using the action selection neural network in accordance with current values of the network parameters to generate an action batch that includes, for each environment replica, a respective action to be performed by the agent in response to the current observation characterizing the current state of the environment replica, wherein the processing comprises executing the inference subgraph to perform batched inference in parallel for the action selection neural network on the current observations in the single tensor to generate a respective network output for each current observation and select a respective action from the respective network output;
generating a transition tuple batch comprising a respective transition tuple for each of the environment replicas, wherein the transition tuple batch is generated by available processors of a plurality of first processors and using the separate processes in parallel, wherein the generating the transition tuple batch comprises constraining each respective sperate process to progress each respective environment replica with a single transition step using the respective action specified in the action batch, the respective transition tuple for each environment replica comprising:
(i) a subsequent observation characterizing a subsequent state that the environment replica transitioned into as a result of the agent performing the respective action in the action batch for the environment replica, and
(ii) a reward generated as a result of the environment replica transitioning into the subsequent state; and
training the action selection neural network by a plurality of second processors, on the transition tuple batch to update the current values of the network parameters using a reinforcement learning technique, the training comprising executing the training subgraph using the plurality of second processors, wherein the training subgraph takes as input the transition tuple batch and the current values of the network parameters and applies the reinforcement learning technique to the training tuples in the batch to generate updated values of the network parameters.
9 . The system of claim 8 , wherein the generating the transition tuple batch comprising the respective transition tuple for each of the environment replicas comprises:
providing the respective action in the action batch to the separate process that maintains the environment replica corresponding to the respective action to cause the environment replica to transition into the subsequent state in parallel; and
obtaining, from each of the separate processes, the subsequent observation and the reward for the environment replica maintained inside of the separate process.
10 . The system of claim 9 , wherein the generating the transition tuple batch further comprises:
after the subsequent observation and the reward have been obtained from all of the separate processes, generating the transition tuple batch from data obtained from the separate processes.
11 . The system of claim 8 , wherein the generating the transition tuple batch comprising the respective transition tuple for each of the environment replicas comprises:
issuing respective calls in parallel to each of the separate processes with the actions for the environment replicas;
waiting until the subsequent observation and the reward are obtained from each of the separate processes in response to the respective calls; and
after determining that a subsequent observation and a reward have been obtained from each of the separate processes, generating the transition tuple batch using the obtained subsequent observations and rewards.
12 . The system of claim 8 , wherein the reinforcement learning technique is a proximal policy optimization (PPO) algorithm.
13 . One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for training an action selection neural network to select actions to be performed by an agent interacting with an environment, wherein the action selection neural network has a plurality of network parameters and is configured to receive an input observation and to process the input observation in accordance with the network parameters to generate a network output that defines an action to be performed by the agent in response to the input observation, wherein the training is performed by processing a computation graph comprising nodes representing neural network operations and edges representing flow of input and output data of the neural network operations, wherein the computation graph includes an inference subgraph and a training subgraph, and wherein the operations comprise, for each of a plurality of steps:
obtaining a plurality of current observations, each current observation characterizing a current state of a respective one of a plurality of environment replicas, wherein each environment replica is maintained by a separate process;
generating an observation batch as a single tensor that combines the plurality of current observations of the plurality of environment replicas maintained by the separate processes;
processing the current observations of the plurality of environment replicas in the single tensor in parallel using the action selection neural network in accordance with current values of the network parameters to generate an action batch that includes, for each environment replica, a respective action to be performed by the agent in response to the current observation characterizing the current state of the environment replica, wherein the processing comprises executing the inference subgraph to perform batched inference in parallel for the action selection neural network on the current observations in the single tensor to generate a respective network output for each current observation and select a respective action from the respective network output;
generating a transition tuple batch comprising a respective transition tuple for each of the environment replicas, wherein the transition tuple batch is generated by available processors of a plurality of first processors and using the separate processes in parallel, wherein the generating the transition tuple batch comprises constraining each respective sperate process to progress each respective environment replica with a single transition step using the respective action specified in the action batch, the respective transition tuple for each environment replica comprising:
(i) a subsequent observation characterizing a subsequent state that the environment replica transitioned into as a result of the agent performing the respective action in the action batch for the environment replica, and
(ii) a reward generated as a result of the environment replica transitioning into the subsequent state; and
training the action selection neural network by a plurality of second processors, on the transition tuple batch to update the current values of the network parameters using a reinforcement learning technique, the training comprising executing the training subgraph using the plurality of second processors, wherein the training subgraph takes as input the transition tuple batch and the current values of the network parameters and applies the reinforcement learning technique to the training tuples in the batch to generate updated values of the network parameters.
14 . The non-transitory computer-readable storage media of claim 13 , wherein the generating the transition tuple batch comprising the respective transition tuple for each of the environment replicas comprises:
providing the respective action in the action batch to the separate process that maintains the environment replica corresponding to the respective action to cause the environment replica to transition into the subsequent state in parallel; and
obtaining, from each of the separate processes, the subsequent observation and the reward for the environment replica maintained inside of the separate process.
15 . The non-transitory computer-readable storage media of claim 14 , wherein the generating the transition tuple batch further comprises:
after the subsequent observation and the reward have been obtained from all of the separate processes, generating the transition tuple batch from data obtained from the separate processes.