修改了算法代码,并建立了一个简单的训练脚本.修改bert处理二维输入,移除PPO的permute参数

Algorithm/__init__.py

@@ -0,0 +1,27 @@
+"""
+MAGAIL Algorithm Package
+
+多智能体生成对抗模仿学习算法实现
+"""
+
+from .magail import MAGAIL
+from .ppo import PPO
+from .disc import GAILDiscrim
+from .bert import Bert
+from .policy import StateIndependentPolicy
+from .buffer import RolloutBuffer
+from .utils import Normalizer, build_mlp, reparameterize, evaluate_lop_pi
+
+__all__ = [
+    'MAGAIL',
+    'PPO',
+    'GAILDiscrim',
+    'Bert',
+    'StateIndependentPolicy',
+    'RolloutBuffer',
+    'Normalizer',
+    'build_mlp',
+    'reparameterize',
+    'evaluate_lop_pi',
+]
+

Algorithm/bert.py

@@ -28,17 +28,26 @@ class Bert(nn.Module):
         self.classifier.train()
 
     def forward(self, x, mask=None):
+        # x可以是2D (batch_size, input_dim) 或 3D (batch_size, seq_len, feature_dim)
+        is_2d_input = (x.dim() == 2)
+        
+        if is_2d_input:
+            # 如果输入是2D,添加一个序列维度
+            x = x.unsqueeze(1)  # (batch_size, 1, input_dim)
+        
         # x: (batch_size, seq_len, input_dim)
         # 线性投影
-        x = self.projection(x)  # (batch_size, input_dim, embed_dim)
+        x = self.projection(x)  # (batch_size, seq_len, embed_dim)
 
         batch_size = x.size(0)
         if self.CLS:
             cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-            x = torch.cat([cls_tokens, x], dim=1)  # (batch_size, 29, embed_dim)
+            x = torch.cat([cls_tokens, x], dim=1)  # (batch_size, seq_len+1, embed_dim)
 
-        # 添加位置编码
-        x = x + self.pos_embed
+        # 添加位置编码(截断或扩展以匹配序列长度)
+        seq_len = x.size(1)
+        pos_embed = self.pos_embed[:, :seq_len, :]
+        x = x + pos_embed
 
         # 转置为(seq_len, batch_size, embed_dim)
         x = x.permute(1, 0, 2)

Algorithm/disc.py

@@ -1,6 +1,9 @@
 import torch
 from torch import nn
-from .bert import Bert
+try:
+    from .bert import Bert
+except ImportError:
+    from bert import Bert
 
 
 DISC_LOGIT_INIT_SCALE = 1.0

Algorithm/magail.py

@@ -2,21 +2,30 @@ import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from .disc import GAILDiscrim
-from .ppo import PPO
-from .utils import Normalizer
+try:
+    from .disc import GAILDiscrim
+    from .ppo import PPO
+    from .utils import Normalizer
+except ImportError:
+    from disc import GAILDiscrim
+    from ppo import PPO
+    from utils import Normalizer
 
 
 class MAGAIL(PPO):
-    def __init__(self, buffer_exp, input_dim, device,
+    def __init__(self, buffer_exp, input_dim, device, action_shape=(2,),
                  disc_coef=20.0, disc_grad_penalty=0.1, disc_logit_reg=0.25, disc_weight_decay=0.0005,
-                 lr_disc=1e-3, epoch_disc=50, batch_size=1000, use_gail_norm=True
+                 lr_disc=1e-3, epoch_disc=50, batch_size=1000, use_gail_norm=True,
+                 **kwargs  # 接受其他PPO参数
                  ):
-        super().__init__(state_shape=input_dim, device=device)
+        super().__init__(state_shape=input_dim, device=device, action_shape=action_shape, **kwargs)
         self.learning_steps = 0
         self.learning_steps_disc = 0
 
-        self.disc = GAILDiscrim(input_dim=input_dim)
+        # 如果input_dim是元组，提取第一个元素
+        state_dim = input_dim[0] if isinstance(input_dim, tuple) else input_dim
+        # 判别器输入是state+next_state拼接，所以维度是state_dim*2
+        self.disc = GAILDiscrim(input_dim=state_dim*2).to(device)  # 移动到指定设备
         self.disc_grad_penalty = disc_grad_penalty
         self.disc_coef = disc_coef
         self.disc_logit_reg = disc_logit_reg
@@ -27,7 +36,9 @@ class MAGAIL(PPO):
 
         self.normalizer = None
         if use_gail_norm:
-            self.normalizer = Normalizer(self.state_shape[0]*2)
+            # state_shape已经是元组形式
+            state_dim = self.state_shape[0] if isinstance(self.state_shape, tuple) else self.state_shape
+            self.normalizer = Normalizer(state_dim*2)
 
         self.batch_size = batch_size
         self.buffer_exp = buffer_exp
@@ -52,7 +63,7 @@ class MAGAIL(PPO):
         # grad penalty
         sample_expert = states_exp_cp
         sample_expert.requires_grad = True
-        disc = self.disc.linear(self.disc.trunk(sample_expert))
+        disc = self.disc(sample_expert)  # 直接调用forward方法
         ones = torch.ones(disc.size(), device=disc.device)
         disc_demo_grad = torch.autograd.grad(disc, sample_expert,
                                              grad_outputs=ones,
@@ -91,7 +102,8 @@ class MAGAIL(PPO):
 
             # Samples from current policy trajectories.
             samples_policy = self.buffer.sample(self.batch_size)
-            states, next_states = samples_policy[1], samples_policy[-3]
+            # samples_policy返回: (states, actions, rewards, dones, tm_dones, log_pis, next_states, means, stds)
+            states, next_states = samples_policy[0], samples_policy[6]  # 修正: 使用states而不是actions
             states = torch.cat([states, next_states], dim=-1)
 
             # Samples from expert demonstrations.
@@ -129,6 +141,8 @@ class MAGAIL(PPO):
         return rewards_t.mean().item() + rewards_i.mean().item()
 
     def save_models(self, path):
+        # 确保目录存在
+        os.makedirs(path, exist_ok=True)
         torch.save({
             'actor': self.actor.state_dict(),
             'critic': self.critic.state_dict(),

Algorithm/policy.py

@@ -1,7 +1,10 @@
 import torch
 import numpy as np
 from torch import nn
-from .utils import build_mlp, reparameterize, evaluate_lop_pi
+try:
+    from .utils import build_mlp, reparameterize, evaluate_lop_pi
+except ImportError:
+    from utils import build_mlp, reparameterize, evaluate_lop_pi
 
 class StateIndependentPolicy(nn.Module):
 

Algorithm/ppo.py

@@ -3,9 +3,14 @@ import torch
 import numpy as np
 from torch import nn
 from torch.optim import Adam
-from buffer import RolloutBuffer
-from bert import Bert
-from policy import StateIndependentPolicy
+try:
+    from .buffer import RolloutBuffer
+    from .bert import Bert
+    from .policy import StateIndependentPolicy
+except ImportError:
+    from buffer import RolloutBuffer
+    from bert import Bert
+    from policy import StateIndependentPolicy
 from abc import ABC, abstractmethod
 
 
@@ -55,7 +60,7 @@ class Algorithm(ABC):
 
 class PPO(Algorithm):
 
-    def __init__(self, state_shape, device, gamma=0.995, rollout_length=2048,
+    def __init__(self, state_shape, device, action_shape=(2,), gamma=0.995, rollout_length=2048,
                  units_actor=(64, 64), epoch_ppo=10, clip_eps=0.2,
                  lambd=0.97, max_grad_norm=1.0, desired_kl=0.01, surrogate_loss_coef=2.,
                  value_loss_coef=5., entropy_coef=0., bounds_loss_coef=10., lr_actor=1e-3, lr_critic=1e-3,
@@ -66,6 +71,7 @@ class PPO(Algorithm):
         self.lr_critic = lr_critic
         self.lr_disc = lr_disc
         self.auto_lr = auto_lr
+        self.action_shape = action_shape
 
         self.use_adv_norm = use_adv_norm
 
@@ -86,8 +92,10 @@ class PPO(Algorithm):
         ).to(device)
 
         # Critic.
+        # 如果state_shape是元组，提取第一个元素
+        state_dim = state_shape[0] if isinstance(state_shape, tuple) else state_shape
         self.critic = Bert(
-            input_dim=state_shape,
+            input_dim=state_dim,
             output_dim=1
         ).to(device)
 
@@ -145,14 +153,12 @@ class PPO(Algorithm):
         targets, gaes = self.calculate_gae(
             values, rewards, dones, tm_dones, next_values, self.gamma, self.lambd)
 
-        state_list = states.permute(1, 0, 2)
-        action_list = actions.permute(1, 0, 2)
-
+        # 处理批量数据（不需要按智能体分组，因为buffer中已经混合了所有智能体的数据）
         for i in range(self.epoch_ppo):
             self.learning_steps_ppo += 1
             self.update_critic(states, targets, writer)
-            for state, action, log_pi in state_list, action_list, log_pi_list:
-                self.update_actor(state, action, log_pi, gaes, mus, sigmas, writer)
+            # 直接使用整个batch进行actor更新
+            self.update_actor(states, actions, log_pi_list, gaes, mus, sigmas, writer)
 
         # self.lr_decay(total_steps, writer)