Systems and methods for end-to-end multi-agent reinforcement learning on a graphics processing unit
Embodiments provide a fast multi-agent reinforcement learning (RL) pipeline that runs the full RL workflow end-to-end on a single GPU, using a single store of data for simulation roll-outs, inference, and training. Specifically, simulations and agents in each simulation are run in tandem, taking advantage of the parallel capabilities of the GPU. This way, the costly GPU-CPU communication and copying is significantly reduced, and simulation sampling and learning rates are in turn improved. In this way, a large number of simulations may be concurrently run on the GPU, thus largely improving efficiency of the RL training.
1 . A graphics processing unit (GPU) based system for implementing multiple reinforcement learning (RL) agents, the GPU based system comprising:
a central processing unit (CPU) host that:
provides, to a GPU device, parameters of a simulation environment involving a plurality of RL agents during a one-time CPU-GPU data transfer at a simulation setup; and
the GPU device communicatively coupled to the CPU host, comprising:
a GPU memory storing a neural network model supporting the plurality of RL agents and a set of simulation roll-out data generated by the plurality of RL agents; and
one or more processing blocks to operate the plurality of RL agents in parallel, generating the set of simulation roll-out data,
wherein each processing block of the one or more processing blocks implements the simulation environment and operates multiple threads in parallel, and
wherein each thread of the multiple threads conducts, via the neural network model stored in the GPU memory, at least one RL agent from the plurality of RL agents to generate at least part of the set of simulation roll-out data by observing a state of the simulation environment including other RL agents, receiving a reward and sampling an action according to a respective policy model corresponding to the at least one RL agent, wherein the set of simulation roll-out data is directly stored, upon generation, at the GPU memory without data copying between the CPU host and the GPU device.
2 . The system of claim 1 , wherein the CPU host transfers data to the GPU device for one time only.
3 . The system of claim 2 , wherein the transferred data includes any combination of:
configuration parameters of the simulation environment;
data arrays for reset; and
initial values for observations, actions, and rewards for reinforcement learning.
4 . The system of claim 2 , wherein the CPU host comprises a data manager that provides an application programming interface (API) to conduct data transfer between the CPU host and the GPU device.
5 . The system of claim 1 , wherein the CPU host comprises a function manager that provides an API to invoke a GPU kernel at the GPU device.
6 . The system of claim 1 , wherein each of the one or more processing blocks executes a replica of the simulation environment in a multi-agent setting.
7 . The system of claim 1 , wherein data arrays generated from the one or more processing blocks during simulation are stored at the GPU memory without transferring the data arrays to the CPU host.
8 . The system of claim 1 , wherein for the at least one thread, the respective policy model outputs probabilities in response to an input of the observed state, and the action is sampled based on the output probabilities.
9 . The system of claim 1 , wherein the roll-out data from the plurality of RL agents and replicas of the simulation environments running on the one or more processing blocks are gathered into a training data batch stored at the GPU memory, and
wherein the training data batch is used for training policy models by the neural network model.
10 . A method for implementing multiple reinforcement learning (RL) agents on a graphics processing unit (GPU) system, the method comprising:
providing, from a central processing unit (CPU) host to a GPU device communicatively coupled to the CPU host, parameters of a simulation environment involving a plurality of RL agents during a one-time CPU-GPU data transfer at a beginning of a simulation,
initializing, by the GPU device, a neural network model stored at a GPU memory to support the plurality of RL agents for the simulation, wherein the GPU memory stores a set of simulation roll-out data generated by the plurality of RL agents;
conducting, at one or more processing blocks of the GPU operated in parallel, wherein at least one processing block of the one or more processing blocks implements the simulation environment and operates multiple threads in parallel, the simulation of the plurality of RL agents including:
generating, via the neural network model, a thread that conducts at least one RL agent from the plurality of RL agent to generate at least part of the set of simulation roll-out data by observing a state of the simulation environment including other RL agents, receiving a reward and sampling an action according to a respective policy model corresponding to the at least one RL agent, wherein the set of simulation roll-out data is directly stored, upon generation, at the GPU memory without data copying between the CPU host and the GPU device,
operating each of the multiple threads by at least one processing block to conduct at least one RL agent from the plurality of RL agents, and
storing, at the GPU memory, simulation data from the multiple threads; and
gathering roll-out data from the simulation for training policy models on the neural network model.
11 . The method of claim 10 , wherein the CPU host transfers data to the GPU device for one time only.
12 . The method of claim 11 , wherein the transferred data includes any combination of:
configuration parameters of the simulation environment;
data arrays for reset;
initial values for observations, actions, and rewards for reinforcement learning.
13 . The method of claim 11 , wherein the CPU host comprises a data manager that provides an application programming interface (API) to conduct data transfer from the CPU host to the GPU device.
14 . The method of claim 10 , wherein the initializing is performed at a function manager in the CPU host that provides an API to invoke a GPU kernel at the GPU device.
15 . The method of claim 10 , wherein each of the one or more processing blocks executes a replica of the simulation environment in a multi-agent setting.
16 . The method of claim 10 , wherein data arrays generated from the one or more processing blocks during simulation are stored at the GPU memory without transferring the data arrays to the CPU host.
17 . The system of claim 10 , wherein for the at least one thread, the respective policy model outputs probabilities in response to an input of the observed state, and the action is sampled based on the output probabilities.
18 . A non-transitory processor-readable medium storing a plurality of processor-executable instructions for implementing multiple reinforcement learning (RL) agents on a graphics processing unit (GPU) system, the instructions being executed by one or more processors to perform operations comprising:
providing, from a central processing unit (CPU) host to a GPU device communicatively coupled to the CPU host, parameters of a simulation environment involving a plurality of RL agents during a one-time CPU-GPU data transfer at a beginning of a simulation,
initializing, by the GPU device, a neural network model stored at a GPU memory to support the plurality of RL agents for the simulation, wherein the GPU memory stores a set of simulation roll-out data generated by the plurality of RL agents;
conducting, at one or more processing blocks of the GPU operated in parallel, wherein at least one processing block of the one or more processing blocks implements the simulation environment and operates multiple threads in parallel, the simulation of the plurality of RL agents including:
generating, via the neural network model, a thread that conducts at least one RL agent from the plurality of RL agent to generate at least part of the set of simulation roll-out data by observing a state of the simulation environment including other RL agents, receiving a reward and sampling an action according to a respective policy model corresponding to the at least one RL agent, wherein the set of simulation roll-out data is directly stored, upon generation, at the GPU memory without data copying between the CPU host and the GPU device,
operating each of the multiple threads by at least one processing block to conduct at least one RL agent from the plurality of RL agents, and
storing, at the GPU memory, simulation data from the multiple threads; and
gathering roll-out data from the simulation for training policy models on the neural network model.
19 . The non-transitory processor-readable medium of claim 18 , wherein the CPU host transfers data to the GPU device for one time only.
20 . The non-transitory processor-readable medium of claim 18 , wherein data arrays generated from the one or more processing blocks during simulation are stored at the GPU memory without transferring the data arrays to the CPU host.