custom_agent.py
MRT2A 上游空间策略网络
原文件 · 229 行 · 9711 字节
third_party/marl_mrt2a/marl_mrt2a/marl/src/modules/agents/custom_agent.py
SHA-256: 8ba538eed4c9ac8a1b00448aceac1b597a06f46990c613ffc49e65549af54930
1from utils.nn_utils import net_from_string23import torch4import torch.nn as nn56import numpy as np7from math import ceil89class CustomAgent(nn.Module):10 def __init__(self, input_shape, args):11 super(CustomAgent, self).__init__()1213 def forward(self, input):14 raise NotImplementedError("Custom agent not implemented")151617class CNNAgent(CustomAgent):18 """把局部栅格观测转换成每个栅格位置的动作分数。"""19 def __init__(self, input_shape, args):20 super().__init__(input_shape, args)2122 self.device = torch.device(args.device)23 self.in_shape = input_shape24 self.in_channels = input_shape[0]25 self.intention = None26 self.current_target_factor = args.current_target_factor27 self.reeval_prob = 1. if (args.agent_reeval_rate is None or args.agent_reeval_rate is True) \28 else (1-np.exp(-args.agent_reeval_rate))29 # reeval_prob 决定已经有目标时是否重新选择目标。30 31 self.n_actions = args.n_actions32 self.n_agents = args.n_agents3334 self.agent_distance_exp = self.dist_grid(input_shape[1], gamma=args.agent_distance_exp)["0"].to(self.device)3536 self.n_out_channels = int(args.n_actions/(input_shape[1]*input_shape[2]))37 self.target_shape = (self.n_out_channels,*input_shape[1:])38 # assert args.action_grid, "Only action grid supported"39 assert self.n_out_channels == 1, "Comms not supported"40 self.net, self.out_shape = net_from_string(args.agent_arch, self.in_shape, target_shape=self.target_shape)41 # 网络输出保持二维布局，每个格子对应一个动作。42 self.net = self.net.to(self.device)4344 # self.dist_given_act = self.dist_grid(self.out_shape[-1], gamma=.5)45 # self.multiplier_act = torch.cat([x.unsqueeze(0) for x in self.dist_given_act.values()],0).to(self.device)4647 def forward(self, input, hidden_state=None, env_info=None):48 # net 保持空间布局输出分数，距离权重会降低远处位置的分数。49 if input.dim() == 3:50 input = input.unsqueeze(0)51 elif input.dim()<3:52 raise ValueError("Not enough dimensions in input")5354 v = self.net(input)55 v = self.agent_distance_exp*v # Apply distance factor56 # intention 保留二维分数图，后续记录和可视化可以直接使用。57 self.intention = v58 try:59 v = torch.flatten(v, start_dim=-3)60 except:61 print(v)6263 v, target_update = self.target_update(v, env_info)64 env_info = self._flatten_env_info(env_info, bs=(v.shape[0]//self.n_agents))65 return v, torch.Tensor([0]), target_update, env_info6667 def init_hidden(self):68 return torch.Tensor([0])6970 def _get_output_shape(self):71 with torch.no_grad():72 v = torch.randn(self.in_shape).unsqueeze(0).to(self.device)73 out = self.net(v)[0]74 return out.shape7576 # v = torch.randn(self.in_shape)77 # v, _ = self.forward(v)78 # return v.shape7980 def target_update(self, v, env_info=None):81 """根据每个智能体当前目标，决定本步是否重新选择目标。"""82 L = v.shape[0]83 v = v.view(-1,self.n_agents,self.n_actions)84 bs = v.shape[0]85 target_update = torch.ones(bs, self.n_agents, 2, dtype=torch.uint8) # 0: revision, 1: change86 if not isinstance(env_info, list):87 env_info = [env_info]*bs888990 for i,worker_env_info in enumerate(env_info):91 if worker_env_info in [None,0] or worker_env_info.get("robot_info", None) is None:92 # 没有当前目标信息时，允许正常选择所有可用动作。93 target_update[i,:,0] = 094 continue # target_update=195 worker_env_info = worker_env_info["robot_info"]96 assert len(worker_env_info)==self.n_agents, "Expected len(env_info)==n_agents, from each worker"97 for j,(pos,cmd) in enumerate(worker_env_info):98 if cmd is None:99 target_update[i,j,0] = 0100 continue101 dif = (cmd[0]-pos[0], cmd[1]-pos[1])102 # current_dif 是当前目标相对智能体中心格的位置。103 v[i,j,:], target_update[i,j,:] = self.target_update_policy(v[i,j,:], dif)104105 v = v.view(L, self.n_actions)106 target_update = target_update.view(L, 2)107 return v, target_update108109 def target_update_policy(self, actions, current_dif):110 """保留当前目标，或调整当前目标分数后重新采样。"""111 if not self._dif_within_obs(current_dif):112 # Continue normally113 return actions, torch.ones((2,))114 115 reeval = True if self.reeval_prob==1 else (np.random.rand()<self.reeval_prob)116 # 不重新选择时，只开放当前目标对应的一个动作。117 if (not reeval) or (self.current_target_factor is None):118 # 只保留当前目标位置，动作选择器会继续执行这个目标。119 actions = -1e10*torch.ones_like(actions)120 actions[self._dif_to_flat(current_dif)] = 1e10121 return actions, torch.zeros((2,))122 else:123 # # Add current target factor124 # actions[self._dif_to_flat(current_dif)] += np.log(self.current_target_factor)125 # Multiply by target factor126 actions[self._dif_to_flat(current_dif)] *= self.current_target_factor127 # 返回值第二项的两列由 runner 记录是否重新选择、是否改变目标。128 return actions, torch.ones((2,))129130 def _dif_within_obs(self, dif):131 sz = self.out_shape[1:]132 return 0 <= dif[0]+sz[0]//2 < sz[0] and 0 <= dif[1]+sz[1]//2 < sz[1]133 134 def _clip_to_obs_range(self, dif):135 sz = self.out_shape[1:]136 bound = (sz[0]//2, sz[1]//2)137 return (min(max(-bound[0], dif[0]), bound[0]-1), min(max(-bound[1], dif[1]), bound[1]-1))138 139 def _dif_to_flat(self, dif):140 """把以智能体为中心的二维偏移转换成一维动作编号。"""141 sz = self.out_shape[1:]142 dif = (dif[0]+sz[0]//2, dif[1]+sz[1]//2)143 try:144 return np.ravel_multi_index(dif, sz)145 except:146 return -1147 148 def _flatten_env_info(self, env_info, bs):149 if not isinstance(env_info, list):150 env_info = [env_info]*bs151 info_list = []152 for i,worker_env_info in enumerate(env_info):153 # worker_info_list = [[-1 for _ in range(self.n_out_channels)]]*self.n_agents154 worker_info_list = [-1]*self.n_agents155 if worker_env_info in [None,0] or worker_env_info.get("robot_info", None) is None:156 # If there's no information, take all actions157 info_list.append(worker_info_list)158 continue159 worker_env_info = worker_env_info["robot_info"]160 assert len(worker_env_info)==self.n_agents, "Expected len(env_info)==n_agents, from each worker"161 for j,(pos,cmd) in enumerate(worker_env_info):162 if cmd is None:163 continue164 dif = (cmd[0]-pos[0], cmd[1]-pos[1])165 # worker_info_list[j] = self._dif_to_flat(dif) if self._dif_within_obs(dif) else [-1 for _ in range(self.n_out_channels)]166 worker_info_list[j] = self._dif_to_flat(dif) if self._dif_within_obs(dif) else -1167 info_list.append(worker_info_list)168 info_list = {"act_info": info_list}169 return info_list170 171 def _flat_to_dif(self, act):172 sz = self.out_shape[1:]173 dif = np.unravel_index(act, sz)174 dif = (dif[0]-sz[0]//2, dif[1]-sz[1]//2)175 return dif176 177 def grid_to_act(self, grid):178 o = self.multiplier_act * grid179 # Sanity check180 o = torch.sum(o, dim=(-1,-2))181 return o182183 def dist_grid(self,sz,gamma=.5):184 """生成距离权重；离中心越远，权重越小。"""185 assert sz%2==1, "Only odd grid sizes supported when converting to action"186 sz = sz+2187 small_sz = ceil(sz/2)188 small_grid = np.full(shape=(small_sz,small_sz), fill_value=small_sz-1)189 for i in range(small_sz-2,-1,-1):190 small_grid[:i+1,:i+1] = i191 large_grid = np.block([[np.flip(small_grid[1:,1:]),np.flipud(small_grid[1:,:])],192 [np.fliplr(small_grid[:,1:]),small_grid]])193 # 从中心、直线移动和斜向移动三种相对位置生成九类距离模板。194 195 c_grid = gamma**large_grid[1:-1,1:-1]196 u_grid = gamma**large_grid[1:-1,2:]197 ur_grid = gamma**large_grid[:-2,2:]198199 dist_given_act = {200 "0": torch.Tensor( c_grid ),201 "U": torch.Tensor( u_grid ),202 "UR": torch.Tensor( ur_grid ),203 "R": torch.Tensor( np.rot90(u_grid).copy() ),204 "DR": torch.Tensor( np.rot90(ur_grid).copy() ),205 "D": torch.Tensor( np.fliplr(u_grid).copy() ),206 "DL": torch.Tensor( np.flip(ur_grid).copy() ),207 "L": torch.Tensor( np.rot90(u_grid, k=-1).copy() ),208 "UL": torch.Tensor( np.rot90(ur_grid, k=-1).copy() ),209 }210 return dist_given_act211212213def xy_print(g):214 for y in range(g.shape[1]):215 print("[", end=" ")216 for x in range(g.shape[0]):217 print(g[x,y], end=" ")218 print(" ]")219220if __name__=="__main__":221 222 ag = CNNAgent((6,9,9), 1)223 v = torch.rand(ag.in_shape)224 act = ag.forward(v)225226 # for k,v in dist_given_act.items():227 # print(k)228 # xy_print(v)229 # print("")