You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardexpand all lines: lightning_examples/warp-drive/multi_agent_rl.py
+45-154
Original file line number
Diff line number
Diff line change
@@ -6,34 +6,46 @@
6
6
# ## Introduction
7
7
8
8
# %% [markdown]
9
-
# This tutorial provides a demonstration of a multi-agent Reinforcement Learning (RL) training loop with [WarpDrive](https://github.com/salesforce/warp-drive). WarpDrive is a flexible, lightweight, and easy-to-use RL framework that implements end-to-end deep multi-agent RL on a single GPU (Graphics Processing Unit). Using the extreme parallelization capability of GPUs, it enables [orders-of-magnitude faster RL](https://arxiv.org/abs/2108.13976) compared to common implementations that blend CPU simulations and GPU models. WarpDrive is extremely efficient as it runs simulations across multiple agents and multiple environment replicas in parallel and completely eliminates the back-and-forth data copying between the CPU and the GPU.
9
+
# This tutorial provides a demonstration of a multi-agent Reinforcement Learning (RL) training loop with [WarpDrive](https://github.com/salesforce/warp-drive). WarpDrive is a flexible, lightweight, and easy-to-use RL framework that implements end-to-end deep multi-agent RL on a GPU (Graphics Processing Unit). Using the extreme parallelization capability of GPUs, it enables [orders-of-magnitude faster RL](https://arxiv.org/abs/2108.13976) compared to common implementations that blend CPU simulations and GPU models. WarpDrive is extremely efficient as it runs simulations across multiple agents and multiple environment replicas all in parallel and completely eliminates the back-and-forth data copying between the CPU and the GPU during every step. As such, WarpDrive
10
+
# - Can simulate 1000s of agents in each environment and thousands of environments in parallel, harnessing the extreme parallelism capability of GPUs.
11
+
# - Eliminates communication between CPU and GPU, and also within the GPU, as read and write operations occur in-place.
12
+
# - Is fully compatible with Pytorch, a highly flexible and very fast deep learning framework.
13
+
# - Implements parallel action sampling on CUDA C, which is ~3x faster than using Pytorch’s sampling methods.
14
+
# - Allows for large-scale distributed training on multiple GPUs.
10
15
#
11
-
# We have integrated WarpDrive with the [Pytorch Lightning](https://www.pytorchlightning.ai/) framework, which greatly reduces the trainer boilerplate code, and improves training flexibility.
16
+
# Below is an overview of WarpDrive’s layout of computational and data structures on a single GPU.
# Computations are organized into blocks, with multiple threads in each block. Each block runs a simulation environment and each thread
19
+
# simulates an agent in an environment. Blocks can access the shared GPU memory that stores simulation data and neural network policy models. A DataManager and FunctionManager enable defining multi-agent RL GPU-workflows with Python APIs. For more details, please read out white [paper](https://arxiv.org/abs/2108.13976).
12
20
#
13
-
# Below, we demonstrate how to use WarpDrive and PytorchLightning together to train a game of [Tag](https://github.com/salesforce/warp-drive/blob/master/example_envs/tag_continuous/tag_continuous.py) where multiple *tagger* agents are trying to run after and tag multiple other *runner* agents. As such, the Warpdrive framework comprises several utility functions that help easily implement any (OpenAI-)*gym-style* RL environment, and furthermore, provides quality-of-life tools to train it end-to-end using just a few lines of code. You may familiarize yourself with WarpDrive with the help of these [tutorials](https://github.com/salesforce/warp-drive/tree/master/tutorials).
21
+
# The Warpdrive framework comprises several utility functions that help easily implement any (OpenAI-)*gym-style* RL environment, and furthermore, provides quality-of-life tools to train it end-to-end using just a few lines of code. You may familiarize yourself with WarpDrive with the help of these [tutorials](https://github.com/salesforce/warp-drive/tree/master/tutorials).
14
22
#
15
23
# We invite everyone to **contribute to WarpDrive**, including adding new multi-agent environments, proposing new features and reporting issues on our open source [repository](https://github.com/salesforce/warp-drive).
24
+
#
25
+
# We have integrated WarpDrive with the [Pytorch Lightning](https://www.pytorchlightning.ai/) framework, which greatly reduces the trainer boilerplate code, and improves training modularity and flexibility. It abstracts away most of the engineering pieces of code, so users can focus on research and building models, and iterate on experiments really fast. Pytorch Lightning also provides support for easily running the model on any hardware, performing distributed training, model checkpointing, performance profiling, logging and visualization.
26
+
#
27
+
# Below, we demonstrate how to use WarpDrive and PytorchLightning together to train a game of [Tag](https://github.com/salesforce/warp-drive/blob/master/example_envs/tag_continuous/tag_continuous.py) where multiple *tagger* agents are trying to run after and tag multiple other *runner* agents. Here's a sample depiction of the game of Tag with $100$ runners and $5$ taggers.
# Uncomment below for enabling animation visualizations.
43
+
# from example_envs.utils.generate_rollout_animation import generate_tag_env_rollout_animation
44
+
# from IPython.display import HTML
45
+
33
46
34
47
# %%
35
-
_NUM_AVAILABLE_GPUS=torch.cuda.device_count()
36
-
assert_NUM_AVAILABLE_GPUS>0, "This notebook needs a GPU to run!"
48
+
asserttorch.cuda.device_count() >0, "This notebook only runs on a GPU!"
37
49
38
50
# %%
39
51
# Set logger level e.g., DEBUG, INFO, WARNING, ERROR.
@@ -46,7 +58,7 @@
46
58
#
47
59
# For our experiment, we consider an environment wherein $5$ taggers and $100$ runners play the game of [Tag](https://github.com/salesforce/warp-drive/blob/master/example_envs/tag_continuous/tag_continuous.py) on a $20 \times 20$ plane. The game lasts $200$ timesteps. Each agent chooses it's own acceleration and turn actions at every timestep, and we use mechanics to determine how the agents move over the grid. When a tagger gets close to a runner, the runner is tagged, and is eliminated from the game. For the configuration below, the runners and taggers have the same unit skill levels, or top speeds.
48
60
#
49
-
# We train the agents using $50$ environments or simulations running in parallel. With WarpDrive, each simulation runs on sepate GPU blocks.
61
+
# We train the agents using $50$ environments or simulations running in parallel. With WarpDrive, each simulation runs on separate GPU blocks.
50
62
#
51
63
# There are two separate policy networks used for the tagger and runner agents. Each network is a fully-connected model with two layers each of $256$ dimensions. We use the Advantage Actor Critic (A2C) algorithm for training. WarpDrive also currently provides the option to use the Proximal Policy Optimization (PPO) algorithm instead.
52
64
@@ -67,10 +79,10 @@
67
79
max_acceleration=0.1,
68
80
# minimum acceleration
69
81
min_acceleration=-0.1,
70
-
# 3*pi/4 radians
71
-
max_turn=2.35,
72
-
# -3*pi/4 radians
73
-
min_turn=-2.35,
82
+
# maximum turn (in radians)
83
+
max_turn=2.35,# 3pi/4 radians
84
+
# minimum turn (in radians)
85
+
min_turn=-2.35,# -3pi/4 radians
74
86
# number of discretized accelerate actions
75
87
num_acceleration_levels=10,
76
88
# number of discretized turn actions
@@ -79,19 +91,21 @@
79
91
skill_level_tagger=1.0,
80
92
# skill level for the runner
81
93
skill_level_runner=1.0,
82
-
# each agent only sees full or partial information
94
+
# each agent sees the full (or partial) information of the world
83
95
use_full_observation=False,
84
96
# flag to indicate if a runner stays in the game after getting tagged
85
97
runner_exits_game_after_tagged=True,
86
98
# number of other agents each agent can see
99
+
# used in the case use_full_observation is False
87
100
num_other_agents_observed=10,
88
-
# positive reward for the tagger upon tagging a runner
101
+
# positive reward for a tagger upon tagging a runner
89
102
tag_reward_for_tagger=10.0,
90
-
# negative reward for the runner upon getting tagged
103
+
# negative reward for a runner upon getting tagged
91
104
tag_penalty_for_runner=-10.0,
92
105
# reward at the end of the game for a runner that isn't tagged
93
106
end_of_game_reward_for_runner=1.0,
94
-
# margin between a tagger and runner to consider the runner as 'tagged'.
107
+
# distance margin between a tagger and runner
108
+
# to consider the runner as being 'tagged'
95
109
tagging_distance=0.02,
96
110
),
97
111
# Trainer settings.
@@ -100,8 +114,9 @@
100
114
num_envs=50,
101
115
# total batch size used for training per iteration (across all the environments)
102
116
train_batch_size=10000,
103
-
# total number of episodes to run the training for (can be arbitrarily high!)
104
-
num_episodes=50000,
117
+
# total number of episodes to run the training for
118
+
# This can be set arbitrarily high!
119
+
num_episodes=500,
105
120
),
106
121
# Policy network settings.
107
122
policy=dict(
@@ -173,136 +188,12 @@
173
188
#
174
189
# We have created a helper function (see below) to visualize an episode rollout. Internally, this function uses the WarpDrive module's `fetch_episode_states` API to fetch the data arrays on the GPU for the duration of an entire episode. Specifically, we fetch the state arrays pertaining to agents' x and y locations on the plane and indicators on which agents are still active in the game. Note that this function may be invoked at any time during training, and it will use the state of the policy models at that time to sample actions and generate the visualization.
# The animation below shows a sample realization of the game episode before training, i.e., with randomly chosen agent actions. The $5$ taggers are marked in pink, while the $100$ blue agents are the runners. Both the taggers and runners move around randomly and about half the runners remain at the end of the episode.
# Note: In the configuration above, we have set the trainer to only train on $50000$ rollout episodes, but you can increase the `num_episodes` configuration parameter to train further. As more training happens, the runners learn to escape the taggers, and the taggers learn to chase after the runner. Sometimes, the taggers also collaborate to team-tag runners. A good number of episodes to train on (for the configuration we have used) is $2$M or higher.
254
+
# Note: In the configuration above, we have set the trainer to only train on $500$ rollout episodes, but you can increase the `num_episodes` configuration parameter to train further. As more training happens, the runners learn to escape the taggers, and the taggers learn to chase after the runner. Sometimes, the taggers also collaborate to team-tag runners. A good number of episodes to train on (for the configuration we have used) is $2$M or higher.
364
255
365
256
# %%
366
257
# Finally, close the WarpDrive module to clear up the CUDA memory heap
0 commit comments