File size: 12,022 Bytes
9b19c29 |
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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 |
from typing import Any, Literal, Protocol, Self, TypeVar, cast, overload
import numpy as np
from overrides import override
from tianshou.data import Batch, ReplayBuffer
from tianshou.data.batch import BatchProtocol, IndexType
from tianshou.data.types import ActBatchProtocol, ObsBatchProtocol, RolloutBatchProtocol
from tianshou.policy import BasePolicy
from tianshou.policy.base import TLearningRateScheduler, TrainingStats
try:
from tianshou.env.pettingzoo_env import PettingZooEnv
except ImportError:
PettingZooEnv = None # type: ignore
class MapTrainingStats(TrainingStats):
def __init__(
self,
agent_id_to_stats: dict[str | int, TrainingStats],
train_time_aggregator: Literal["min", "max", "mean"] = "max",
) -> None:
self._agent_id_to_stats = agent_id_to_stats
train_times = [agent_stats.train_time for agent_stats in agent_id_to_stats.values()]
match train_time_aggregator:
case "max":
aggr_function = max
case "min":
aggr_function = min
case "mean":
aggr_function = np.mean # type: ignore
case _:
raise ValueError(
f"Unknown {train_time_aggregator=}",
)
self.train_time = aggr_function(train_times)
self.smoothed_loss = {}
@override
def get_loss_stats_dict(self) -> dict[str, float]:
"""Collects loss_stats_dicts from all agents, prepends agent_id to all keys, and joins results."""
result_dict = {}
for agent_id, stats in self._agent_id_to_stats.items():
agent_loss_stats_dict = stats.get_loss_stats_dict()
for k, v in agent_loss_stats_dict.items():
result_dict[f"{agent_id}/" + k] = v
return result_dict
class MAPRolloutBatchProtocol(RolloutBatchProtocol, Protocol):
# TODO: this might not be entirely correct.
# The whole MAP data processing pipeline needs more documentation and possibly some refactoring
@overload
def __getitem__(self, index: str) -> RolloutBatchProtocol:
...
@overload
def __getitem__(self, index: IndexType) -> Self:
...
def __getitem__(self, index: str | IndexType) -> Any:
...
class MultiAgentPolicyManager(BasePolicy):
"""Multi-agent policy manager for MARL.
This multi-agent policy manager accepts a list of
:class:`~tianshou.policy.BasePolicy`. It dispatches the batch data to each
of these policies when the "forward" is called. The same as "process_fn"
and "learn": it splits the data and feeds them to each policy. A figure in
:ref:`marl_example` can help you better understand this procedure.
:param policies: a list of policies.
:param env: a PettingZooEnv.
:param action_scaling: if True, scale the action from [-1, 1] to the range
of action_space. Only used if the action_space is continuous.
:param action_bound_method: method to bound action to range [-1, 1].
Only used if the action_space is continuous.
:param lr_scheduler: if not None, will be called in `policy.update()`.
"""
def __init__(
self,
*,
policies: list[BasePolicy],
# TODO: 1 why restrict to PettingZooEnv?
# TODO: 2 This is the only policy that takes an env in init, is it really needed?
env: PettingZooEnv,
action_scaling: bool = False,
action_bound_method: Literal["clip", "tanh"] | None = "clip",
lr_scheduler: TLearningRateScheduler | None = None,
) -> None:
super().__init__(
action_space=env.action_space,
observation_space=env.observation_space,
action_scaling=action_scaling,
action_bound_method=action_bound_method,
lr_scheduler=lr_scheduler,
)
assert len(policies) == len(env.agents), "One policy must be assigned for each agent."
self.agent_idx = env.agent_idx
for i, policy in enumerate(policies):
# agent_id 0 is reserved for the environment proxy
# (this MultiAgentPolicyManager)
policy.set_agent_id(env.agents[i])
self.policies: dict[str | int, BasePolicy] = dict(zip(env.agents, policies, strict=True))
"""Maps agent_id to policy."""
# TODO: unused - remove it?
def replace_policy(self, policy: BasePolicy, agent_id: int) -> None:
"""Replace the "agent_id"th policy in this manager."""
policy.set_agent_id(agent_id)
self.policies[agent_id] = policy
# TODO: violates Liskov substitution principle
def process_fn( # type: ignore
self,
batch: MAPRolloutBatchProtocol,
buffer: ReplayBuffer,
indice: np.ndarray,
) -> MAPRolloutBatchProtocol:
"""Dispatch batch data from `obs.agent_id` to every policy's process_fn.
Save original multi-dimensional rew in "save_rew", set rew to the
reward of each agent during their "process_fn", and restore the
original reward afterwards.
"""
# TODO: maybe only str is actually allowed as agent_id? See MAPRolloutBatchProtocol
results: dict[str | int, RolloutBatchProtocol] = {}
assert isinstance(
batch.obs,
BatchProtocol,
), f"here only observations of type Batch are permitted, but got {type(batch.obs)}"
# reward can be empty Batch (after initial reset) or nparray.
has_rew = isinstance(buffer.rew, np.ndarray)
if has_rew: # save the original reward in save_rew
# Since we do not override buffer.__setattr__, here we use _meta to
# change buffer.rew, otherwise buffer.rew = Batch() has no effect.
save_rew, buffer._meta.rew = buffer.rew, Batch() # type: ignore
for agent, policy in self.policies.items():
agent_index = np.nonzero(batch.obs.agent_id == agent)[0]
if len(agent_index) == 0:
results[agent] = cast(RolloutBatchProtocol, Batch())
continue
tmp_batch, tmp_indice = batch[agent_index], indice[agent_index]
if has_rew:
tmp_batch.rew = tmp_batch.rew[:, self.agent_idx[agent]]
buffer._meta.rew = save_rew[:, self.agent_idx[agent]]
if not hasattr(tmp_batch.obs, "mask"):
if hasattr(tmp_batch.obs, "obs"):
tmp_batch.obs = tmp_batch.obs.obs
if hasattr(tmp_batch.obs_next, "obs"):
tmp_batch.obs_next = tmp_batch.obs_next.obs
results[agent] = policy.process_fn(tmp_batch, buffer, tmp_indice)
if has_rew: # restore from save_rew
buffer._meta.rew = save_rew
return Batch(results)
_TArrOrActBatch = TypeVar("_TArrOrActBatch", bound="np.ndarray | ActBatchProtocol")
def exploration_noise(
self,
act: _TArrOrActBatch,
batch: ObsBatchProtocol,
) -> _TArrOrActBatch:
"""Add exploration noise from sub-policy onto act."""
if not isinstance(batch.obs, Batch):
raise TypeError(
f"here only observations of type Batch are permitted, but got {type(batch.obs)}",
)
for agent_id, policy in self.policies.items():
agent_index = np.nonzero(batch.obs.agent_id == agent_id)[0]
if len(agent_index) == 0:
continue
act[agent_index] = policy.exploration_noise(act[agent_index], batch[agent_index])
return act
def forward( # type: ignore
self,
batch: Batch,
state: dict | Batch | None = None,
**kwargs: Any,
) -> Batch:
"""Dispatch batch data from obs.agent_id to every policy's forward.
:param batch: TODO: document what is expected at input and make a BatchProtocol for it
:param state: if None, it means all agents have no state. If not
None, it should contain keys of "agent_1", "agent_2", ...
:return: a Batch with the following contents:
TODO: establish a BatcProtocol for this
::
{
"act": actions corresponding to the input
"state": {
"agent_1": output state of agent_1's policy for the state
"agent_2": xxx
...
"agent_n": xxx}
"out": {
"agent_1": output of agent_1's policy for the input
"agent_2": xxx
...
"agent_n": xxx}
}
"""
results: list[tuple[bool, np.ndarray, Batch, np.ndarray | Batch, Batch]] = []
for agent_id, policy in self.policies.items():
# This part of code is difficult to understand.
# Let's follow an example with two agents
# batch.obs.agent_id is [1, 2, 1, 2, 1, 2] (with batch_size == 6)
# each agent plays for three transitions
# agent_index for agent 1 is [0, 2, 4]
# agent_index for agent 2 is [1, 3, 5]
# we separate the transition of each agent according to agent_id
agent_index = np.nonzero(batch.obs.agent_id == agent_id)[0]
if len(agent_index) == 0:
# (has_data, agent_index, out, act, state)
results.append((False, np.array([-1]), Batch(), Batch(), Batch()))
continue
tmp_batch = batch[agent_index]
if "rew" in tmp_batch.get_keys() and isinstance(tmp_batch.rew, np.ndarray):
# reward can be empty Batch (after initial reset) or nparray.
tmp_batch.rew = tmp_batch.rew[:, self.agent_idx[agent_id]]
if not hasattr(tmp_batch.obs, "mask"):
if hasattr(tmp_batch.obs, "obs"):
tmp_batch.obs = tmp_batch.obs.obs
if hasattr(tmp_batch.obs_next, "obs"):
tmp_batch.obs_next = tmp_batch.obs_next.obs
out = policy(
batch=tmp_batch,
state=None if state is None else state[agent_id],
**kwargs,
)
act = out.act
each_state = out.state if (hasattr(out, "state") and out.state is not None) else Batch()
results.append((True, agent_index, out, act, each_state))
holder: Batch = Batch.cat(
[{"act": act} for (has_data, agent_index, out, act, each_state) in results if has_data],
)
state_dict, out_dict = {}, {}
for (agent_id, _), (has_data, agent_index, out, act, state) in zip(
self.policies.items(),
results,
strict=True,
):
if has_data:
holder.act[agent_index] = act
state_dict[agent_id] = state
out_dict[agent_id] = out
holder["out"] = out_dict
holder["state"] = state_dict
return holder
# Violates Liskov substitution principle
def learn( # type: ignore
self,
batch: MAPRolloutBatchProtocol,
*args: Any,
**kwargs: Any,
) -> MapTrainingStats:
"""Dispatch the data to all policies for learning.
:param batch: must map agent_ids to rollout batches
"""
agent_id_to_stats = {}
for agent_id, policy in self.policies.items():
data = batch[agent_id]
if not data.is_empty():
train_stats = policy.learn(batch=data, **kwargs)
agent_id_to_stats[agent_id] = train_stats
return MapTrainingStats(agent_id_to_stats)
# Need a train method that set all sub-policies to train mode.
# No need for a similar eval function, as eval internally uses the train function.
def train(self, mode: bool = True) -> Self:
"""Set each internal policy in training mode."""
for policy in self.policies.values():
policy.train(mode)
return self
|