使用 Python TorchRL 进行多代理强化学习

!pip3 install torchrl  
 !pip3 install vmas  
 !pip3 install tqdm
1.
2.
3.

import torch  
 from torch import multiprocessing  
   
 # 设置设备  
 is_fork = multiprocessing.get_start_method() == "fork"  
 device = torch.device(0) if torch.cuda.is_available() and not is_fork else torch.device("cpu")  
 vmas_device = device  # 运行 VMAS 模拟器的设备  
 # 采样和训练参数  
 frames_per_batch = 6000    
 n_iters = 10    
 total_frames = frames_per_batch * n_iters  
 # 训练细节  
 num_epochs = 30    
 minibatch_size = 400    
 lr = 3e-4    
 max_grad_norm = 1.0    
 # PPO 参数  
 clip_epsilon = 0.2    
 gamma = 0.99    
 lmbda = 0.9    
 entropy_eps = 1e-4
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.

from torchrl.envs.libs.vmas import VmasEnv  
   
 max_steps = 100    
 num_vmas_envs = frames_per_batch // max_steps    
 scenario_name = "navigation"  
 n_agents = 3  
 env = VmasEnv(  
     scenario=scenario_name,  
     num_envs=num_vmas_envs,  
     continuous_actions=True,  
     max_steps=max_steps,  
     device=vmas_device,  
     n_agents=n_agents,  
 )
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.

from torch.nn import Sequential, Tanh  
 from tensordict.nn import TensorDictModule  
 from torchrl.modules import MultiAgentMLP, ProbabilisticActor, TanhNormal  
 from tensordict.nn.distributions import NormalParamExtractor  
   
 share_parameters_policy = True  
 # 定义策略网络  
 policy_net = Sequential(  
     MultiAgentMLP(  
         n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],  
         n_agent_outputs=2 * env.action_spec.shape[-1],  
         n_agents=env.n_agents,  
         centralised=False,  
         share_params=share_parameters_policy,  
         device=device,  
         depth=2,  
         num_cells=256,  
         activation_class=Tanh,  
    ),  
     NormalParamExtractor(),  
 )  
 # 将网络包装在 TensorDictModule 中  
 policy_module = TensorDictModule(  
     policy_net,  
     in_keys=[("agents", "observation")],  
     out_keys=[("agents", "loc"), ("agents", "scale")],  
 )  
 # 创建概率性行动者  
 policy = ProbabilisticActor(  
     module=policy_module,  
     spec=env.unbatched_action_spec,  
     in_keys=[("agents", "loc"), ("agents", "scale")],  
     out_keys=[env.action_key],  
     distribution_class=TanhNormal,  
     distribution_kwargs={  
         "low": env.unbatched_action_spec[env.action_key].space.low,  
         "high": env.unbatched_action_spec[env.action_key].space.high,  
    },  
     return_log_prob=True,  
     log_prob_key=("agents", "sample_log_prob"),  
 )
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.

share_parameters_critic = True  
 mappo = True  # 设置为 False 以使用 IPPO  
   
 critic_net = MultiAgentMLP(  
     n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],  
     n_agent_outputs=1,    
     n_agents=env.n_agents,  
     centralised=mappo,  
     share_params=share_parameters_critic,  
     device=device,  
     depth=2,  
     num_cells=256,  
     activation_class=Tanh,  
 )  
 critic = TensorDictModule(  
     module=critic_net,  
     in_keys=[("agents", "observation")],  
     out_keys=[("agents", "state_value")],  
 )
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.

from torchrl.collectors import SyncDataCollector  
   
 collector = SyncDataCollector(  
     env,  
     policy,  
     device=vmas_device,  
     storing_device=device,  
     frames_per_batch=frames_per_batch,  
     total_frames=total_frames,  
 )
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

from torchrl.data.replay_buffers import ReplayBuffer  
 from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement  
 from torchrl.data.replay_buffers.storages import LazyTensorStorage  
 from torchrl.objectives import ClipPPOLoss, ValueEstimators  
 from tqdm import tqdm  
   
 # 重放缓冲区设置  
 replay_buffer = ReplayBuffer(  
     storage=LazyTensorStorage(  
         size=frames_per_batch,  
         device=device,  
    ),  
     sampler=SamplerWithoutReplacement(),  
 )  
 # PPO 损失设置  
 ppo_loss = ClipPPOLoss(  
     actor=policy,  
     critic=critic,  
     clip_epsilon=clip_epsilon,  
     entropy_bonus=entropy_eps,  
     value_loss_coef=0.5,  
     gamma=gamma,  
     lam=lmbda,  
     advantage_normalization=True,  
 )  
 # 训练循环  
 for i in tqdm(range(n_iters)):  
     batch = collector.next()  
     replay_buffer.extend(batch)  
     for _ in range(num_epochs):  
         for minibatch in replay_buffer.sample(minibatch_size):  
             loss = ppo_loss(minibatch)  
             loss.backward()  
             torch.nn.utils.clip_grad_norm_(ppo_loss.parameters(), max_grad_norm)  
             optimizer.step()  
             optimizer.zero_grad()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.