ConsensusDecisionMaking/main_test.py at main · shushukunkunkun/ConsensusDecisionMaking · GitHub

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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import numpy as np
import torch
from qmix import QMIX  # 从qmix.py导入QMIX类
from multi_agent_env import MultiAgentEnv  # 从multi_agent_env.py导入MultiAgentEnv类
def main():
    # 定义参数
    class Args:
        n_agents = 2
        n_actions = 27
        state_shape = 14
        obs_shape = 10
        rnn_hidden_dim = 128
        qmix_hidden_dim = 128
        hypernet_hidden_dim = 128
        lr = 0.001
        gamma = 1
        epsilon = 0.1
        target_update_cycle = 200
        grad_norm_clip = 10
        save_cycle = 500
        load_model = False
        map_path = r'C:\Users\pc\Desktop\Python工程文件\RL4ConflictResolution\map\NarrowCorridor_3.png'
        model_dir = './models'
        alg = 'qmix'
        map = 'narrow_corridor'
        cuda = torch.cuda.is_available()
        last_action = False
        reuse_network = False
        two_hyper_layers = False  # 是否使用两层全连接层

    args = Args()
    # 训练设置
    n_episodes = 1000
    max_episode_len = 100  # 单次训练最多步数
    train_step = 0

    # 初始化环境和QMIX网络
    # 初始化环境
    qmix = QMIX(args)
    env = MultiAgentEnv(args.map_path, qmix, visualize_every_frame=False)


    # 初始化批次数据存储
    episodes = []  # 用于存储多个episode的数据

    for episode in range(n_episodes):
        # 重置环境
        qmix.init_hidden(episode_num=args.n_agents)
        state = env.reset()
        episode_data = {
            's': [],  # 全局状态
            'o': [],  # 每个智能体的观测
            'o_next': [],  # 每个智能体的下一个观测
            'u': [],  # 动作
            'r': [],  # 奖励
            'terminated': [],  # 终止标志
            'padded': []  # 填充标志
        }

        for step in range(max_episode_len):
            # 获取每个智能体的动作
            actions = []
            for agent in env.agents:
                agent_state = agent.get_state()
                action = agent.select_action(agent_state, qmix)
                actions.append(action)

            # 执行动作并获取下一状态和奖励
            next_state, rewards, done = env.step(actions)

            # 存储每一步的数据
            episode_data['s'].append(env.get_global_state())
            episode_data['o'].append(np.concatenate(state))  # 将state展开成一维数组
            episode_data['o_next'].append(np.concatenate(next_state))  # 将next_state展开成一维数组
            episode_data['u'].append(actions)
            episode_data['r'].append(rewards)
            episode_data['terminated'].append([done] * args.n_agents)
            episode_data['padded'].append([0] * args.n_agents)  # 填充标志为0，表示未填充

            # 更新状态
            state = next_state

            if done:
                break

        # 如果episode不足最大长度，需要填充数据
        for _ in range(step + 1, max_episode_len):
            episode_data['s'].append(np.zeros_like(episode_data['s'][0]))
            episode_data['o'].append(np.zeros_like(episode_data['o'][0]))
            episode_data['o_next'].append(np.zeros_like(episode_data['o_next'][0]))
            episode_data['u'].append(np.zeros_like(episode_data['u'][0]))
            episode_data['r'].append([0] * args.n_agents)
            episode_data['terminated'].append([1] * args.n_agents)
            episode_data['padded'].append([1] * args.n_agents)  # 填充标志为1，表示填充部分

        # 将单个episode的数据添加到批次数据中
        episodes.append({
            's': np.array(episode_data['s']),
            'o': np.array(episode_data['o']),
            'o_next': np.array(episode_data['o_next']),
            'u': np.array(episode_data['u']),
            'r': np.array(episode_data['r']),
            'terminated': np.array(episode_data['terminated']),
            'padded': np.array(episode_data['padded'])
        })

        # 构建批次数据
        batch = {
            's': np.array([ep['s'] for ep in episodes]),
            'o': np.array([ep['o'] for ep in episodes]),
            'o_next': np.array([ep['o_next'] for ep in episodes]),
            'u': np.array([ep['u'] for ep in episodes]),
            'r': np.array([ep['r'] for ep in episodes]),
            'terminated': np.array([ep['terminated'] for ep in episodes]),
            'padded': np.array([ep['padded'] for ep in episodes])
        }

        # 学习
        qmix.learn(batch, max_episode_len, train_step)
        train_step += 1

        # 每隔一定周期保存模型
        if episode % args.save_cycle == 0:
            qmix.save_model(train_step)
            print(f"Episode {episode}: Model saved")

    print("Training completed.")

if __name__ == "__main__":
    main()