torchrl 0.11.0__cp314-cp314t-win_amd64.whl

sota-implementations/multiagent/mappo_ippo.py

@@ -0,0 +1,267 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import time
+
+import hydra
+import torch
+from tensordict.nn import TensorDictModule
+from tensordict.nn.distributions import NormalParamExtractor
+from torch import nn
+from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import TensorDictReplayBuffer
+from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
+from torchrl.data.replay_buffers.storages import LazyTensorStorage
+from torchrl.envs import RewardSum, TransformedEnv
+from torchrl.envs.libs.vmas import VmasEnv
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import ProbabilisticActor, TanhNormal, ValueOperator
+from torchrl.modules.models.multiagent import MultiAgentMLP
+from torchrl.objectives import ClipPPOLoss, ValueEstimators
+from utils.logging import init_logging, log_evaluation, log_training
+from utils.utils import DoneTransform
+
+
+def rendering_callback(env, td):
+    env.frames.append(env.render(mode="rgb_array", agent_index_focus=None))
+
+
+@hydra.main(version_base="1.1", config_path="", config_name="mappo_ippo")
+def train(cfg: DictConfig):  # noqa: F821
+    # Device
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
+    cfg.env.device = cfg.train.device
+
+    # Seeding
+    torch.manual_seed(cfg.seed)
+
+    # Sampling
+    cfg.env.vmas_envs = cfg.collector.frames_per_batch // cfg.env.max_steps
+    cfg.collector.total_frames = cfg.collector.frames_per_batch * cfg.collector.n_iters
+    cfg.buffer.memory_size = cfg.collector.frames_per_batch
+
+    # Create env and env_test
+    env = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.env.vmas_envs,
+        continuous_actions=True,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+    env = TransformedEnv(
+        env,
+        RewardSum(in_keys=[env.reward_key], out_keys=[("agents", "episode_reward")]),
+    )
+
+    env_test = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.eval.evaluation_episodes,
+        continuous_actions=True,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+
+    # Policy
+    actor_net = nn.Sequential(
+        MultiAgentMLP(
+            n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
+            n_agent_outputs=2
+            * env.full_action_spec_unbatched[env.action_key].shape[-1],
+            n_agents=env.n_agents,
+            centralised=False,
+            share_params=cfg.model.shared_parameters,
+            device=cfg.train.device,
+            depth=2,
+            num_cells=256,
+            activation_class=nn.Tanh,
+        ),
+        NormalParamExtractor(),
+    )
+    policy_module = TensorDictModule(
+        actor_net,
+        in_keys=[("agents", "observation")],
+        out_keys=[("agents", "loc"), ("agents", "scale")],
+    )
+    policy = ProbabilisticActor(
+        module=policy_module,
+        spec=env.full_action_spec_unbatched,
+        in_keys=[("agents", "loc"), ("agents", "scale")],
+        out_keys=[env.action_key],
+        distribution_class=TanhNormal,
+        distribution_kwargs={
+            "low": env.full_action_spec_unbatched[("agents", "action")].space.low,
+            "high": env.full_action_spec_unbatched[("agents", "action")].space.high,
+        },
+        return_log_prob=True,
+    )
+
+    # Critic
+    module = MultiAgentMLP(
+        n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
+        n_agent_outputs=1,
+        n_agents=env.n_agents,
+        centralised=cfg.model.centralised_critic,
+        share_params=cfg.model.shared_parameters,
+        device=cfg.train.device,
+        depth=2,
+        num_cells=256,
+        activation_class=nn.Tanh,
+    )
+    value_module = ValueOperator(
+        module=module,
+        in_keys=[("agents", "observation")],
+    )
+
+    collector = SyncDataCollector(
+        env,
+        policy,
+        device=cfg.env.device,
+        storing_device=cfg.train.device,
+        frames_per_batch=cfg.collector.frames_per_batch,
+        total_frames=cfg.collector.total_frames,
+        postproc=DoneTransform(reward_key=env.reward_key, done_keys=env.done_keys),
+    )
+
+    replay_buffer = TensorDictReplayBuffer(
+        storage=LazyTensorStorage(cfg.buffer.memory_size, device=cfg.train.device),
+        sampler=SamplerWithoutReplacement(),
+        batch_size=cfg.train.minibatch_size,
+    )
+
+    # Loss
+    loss_module = ClipPPOLoss(
+        actor_network=policy,
+        critic_network=value_module,
+        clip_epsilon=cfg.loss.clip_epsilon,
+        entropy_coeff=cfg.loss.entropy_eps,
+        normalize_advantage=False,
+    )
+    loss_module.set_keys(
+        reward=env.reward_key,
+        action=env.action_key,
+        done=("agents", "done"),
+        terminated=("agents", "terminated"),
+    )
+    loss_module.make_value_estimator(
+        ValueEstimators.GAE, gamma=cfg.loss.gamma, lmbda=cfg.loss.lmbda
+    )
+    optim = torch.optim.Adam(loss_module.parameters(), cfg.train.lr)
+
+    # Logging
+    if cfg.logger.backend:
+        model_name = (
+            ("Het" if not cfg.model.shared_parameters else "")
+            + ("MA" if cfg.model.centralised_critic else "I")
+            + "PPO"
+        )
+        logger = init_logging(cfg, model_name)
+
+    total_time = 0
+    total_frames = 0
+    sampling_start = time.time()
+    for i, tensordict_data in enumerate(collector):
+        torchrl_logger.info(f"\nIteration {i}")
+
+        sampling_time = time.time() - sampling_start
+
+        with torch.no_grad():
+            loss_module.value_estimator(
+                tensordict_data,
+                params=loss_module.critic_network_params,
+                target_params=loss_module.target_critic_network_params,
+            )
+        current_frames = tensordict_data.numel()
+        total_frames += current_frames
+        data_view = tensordict_data.reshape(-1)
+        replay_buffer.extend(data_view)
+
+        training_tds = []
+        training_start = time.time()
+        for _ in range(cfg.train.num_epochs):
+            for _ in range(cfg.collector.frames_per_batch // cfg.train.minibatch_size):
+                subdata = replay_buffer.sample()
+                loss_vals = loss_module(subdata)
+                training_tds.append(loss_vals.detach())
+
+                loss_value = (
+                    loss_vals["loss_objective"]
+                    + loss_vals["loss_critic"]
+                    + loss_vals["loss_entropy"]
+                )
+
+                loss_value.backward()
+
+                total_norm = torch.nn.utils.clip_grad_norm_(
+                    loss_module.parameters(), cfg.train.max_grad_norm
+                )
+                training_tds[-1].set("grad_norm", total_norm.mean())
+
+                optim.step()
+                optim.zero_grad()
+
+        collector.update_policy_weights_()
+
+        training_time = time.time() - training_start
+
+        iteration_time = sampling_time + training_time
+        total_time += iteration_time
+        training_tds = torch.stack(training_tds)
+
+        # More logs
+        if cfg.logger.backend:
+            log_training(
+                logger,
+                training_tds,
+                tensordict_data,
+                sampling_time,
+                training_time,
+                total_time,
+                i,
+                current_frames,
+                total_frames,
+                step=i,
+            )
+
+        if (
+            cfg.eval.evaluation_episodes > 0
+            and i % cfg.eval.evaluation_interval == 0
+            and cfg.logger.backend
+        ):
+            evaluation_start = time.time()
+            with torch.no_grad(), set_exploration_type(ExplorationType.DETERMINISTIC):
+                env_test.frames = []
+                rollouts = env_test.rollout(
+                    max_steps=cfg.env.max_steps,
+                    policy=policy,
+                    callback=rendering_callback,
+                    auto_cast_to_device=True,
+                    break_when_any_done=False,
+                    # We are running vectorized evaluation we do not want it to stop when just one env is done
+                )
+
+                evaluation_time = time.time() - evaluation_start
+
+                log_evaluation(logger, rollouts, env_test, evaluation_time, step=i)
+
+        if cfg.logger.backend == "wandb":
+            logger.experiment.log({}, commit=True)
+        sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
+
+
+if __name__ == "__main__":
+    train()

sota-implementations/multiagent/qmix_vdn.py

@@ -0,0 +1,271 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import time
+
+import hydra
+import torch
+from tensordict.nn import TensorDictModule, TensorDictSequential
+from torch import nn
+from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import TensorDictReplayBuffer
+from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
+from torchrl.data.replay_buffers.storages import LazyTensorStorage
+from torchrl.envs import RewardSum, TransformedEnv
+from torchrl.envs.libs.vmas import VmasEnv
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import EGreedyModule, QValueModule, SafeSequential
+from torchrl.modules.models.multiagent import MultiAgentMLP, QMixer, VDNMixer
+from torchrl.objectives import SoftUpdate, ValueEstimators
+from torchrl.objectives.multiagent.qmixer import QMixerLoss
+from utils.logging import init_logging, log_evaluation, log_training
+
+
+def rendering_callback(env, td):
+    env.frames.append(env.render(mode="rgb_array", agent_index_focus=None))
+
+
+@hydra.main(version_base="1.1", config_path="", config_name="qmix_vdn")
+def train(cfg: DictConfig):  # noqa: F821
+    # Device
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
+    cfg.env.device = cfg.train.device
+
+    # Seeding
+    torch.manual_seed(cfg.seed)
+
+    # Sampling
+    cfg.env.vmas_envs = cfg.collector.frames_per_batch // cfg.env.max_steps
+    cfg.collector.total_frames = cfg.collector.frames_per_batch * cfg.collector.n_iters
+    cfg.buffer.memory_size = cfg.collector.frames_per_batch
+
+    # Create env and env_test
+    env = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.env.vmas_envs,
+        continuous_actions=False,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+    env = TransformedEnv(
+        env,
+        RewardSum(in_keys=[env.reward_key], out_keys=[("agents", "episode_reward")]),
+    )
+
+    env_test = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.eval.evaluation_episodes,
+        continuous_actions=False,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+
+    # Policy
+    net = MultiAgentMLP(
+        n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
+        n_agent_outputs=env.full_action_spec["agents", "action"].space.n,
+        n_agents=env.n_agents,
+        centralised=False,
+        share_params=cfg.model.shared_parameters,
+        device=cfg.train.device,
+        depth=2,
+        num_cells=256,
+        activation_class=nn.Tanh,
+    )
+    module = TensorDictModule(
+        net, in_keys=[("agents", "observation")], out_keys=[("agents", "action_value")]
+    )
+    value_module = QValueModule(
+        action_value_key=("agents", "action_value"),
+        out_keys=[
+            env.action_key,
+            ("agents", "action_value"),
+            ("agents", "chosen_action_value"),
+        ],
+        spec=env.full_action_spec_unbatched,
+        action_space=None,
+    )
+    qnet = SafeSequential(module, value_module)
+
+    qnet_explore = TensorDictSequential(
+        qnet,
+        EGreedyModule(
+            eps_init=0.3,
+            eps_end=0,
+            annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
+            action_key=env.action_key,
+            spec=env.full_action_spec_unbatched,
+        ),
+    )
+
+    if cfg.loss.mixer_type == "qmix":
+        mixer = TensorDictModule(
+            module=QMixer(
+                state_shape=env.observation_spec_unbatched[
+                    "agents", "observation"
+                ].shape,
+                mixing_embed_dim=32,
+                n_agents=env.n_agents,
+                device=cfg.train.device,
+            ),
+            in_keys=[("agents", "chosen_action_value"), ("agents", "observation")],
+            out_keys=["chosen_action_value"],
+        )
+    elif cfg.loss.mixer_type == "vdn":
+        mixer = TensorDictModule(
+            module=VDNMixer(
+                n_agents=env.n_agents,
+                device=cfg.train.device,
+            ),
+            in_keys=[("agents", "chosen_action_value")],
+            out_keys=["chosen_action_value"],
+        )
+    else:
+        raise ValueError("Mixer type not in the example")
+
+    collector = SyncDataCollector(
+        env,
+        qnet_explore,
+        device=cfg.env.device,
+        storing_device=cfg.train.device,
+        frames_per_batch=cfg.collector.frames_per_batch,
+        total_frames=cfg.collector.total_frames,
+    )
+
+    replay_buffer = TensorDictReplayBuffer(
+        storage=LazyTensorStorage(cfg.buffer.memory_size, device=cfg.train.device),
+        sampler=SamplerWithoutReplacement(),
+        batch_size=cfg.train.minibatch_size,
+    )
+
+    loss_module = QMixerLoss(qnet, mixer, delay_value=True)
+    loss_module.set_keys(
+        action_value=("agents", "action_value"),
+        local_value=("agents", "chosen_action_value"),
+        global_value="chosen_action_value",
+        action=env.action_key,
+    )
+    loss_module.make_value_estimator(ValueEstimators.TD0, gamma=cfg.loss.gamma)
+    target_net_updater = SoftUpdate(loss_module, eps=1 - cfg.loss.tau)
+
+    optim = torch.optim.Adam(loss_module.parameters(), cfg.train.lr)
+
+    # Logging
+    if cfg.logger.backend:
+        model_name = (
+            "Het" if not cfg.model.shared_parameters else ""
+        ) + cfg.loss.mixer_type.upper()
+        logger = init_logging(cfg, model_name)
+
+    total_time = 0
+    total_frames = 0
+    sampling_start = time.time()
+    for i, tensordict_data in enumerate(collector):
+        torchrl_logger.info(f"\nIteration {i}")
+
+        sampling_time = time.time() - sampling_start
+
+        # Remove agent dimension from reward (since it is shared in QMIX/VDN)
+        tensordict_data.set(
+            ("next", "reward"), tensordict_data.get(("next", env.reward_key)).mean(-2)
+        )
+        del tensordict_data["next", env.reward_key]
+        tensordict_data.set(
+            ("next", "episode_reward"),
+            tensordict_data.get(("next", "agents", "episode_reward")).mean(-2),
+        )
+        del tensordict_data["next", "agents", "episode_reward"]
+
+        current_frames = tensordict_data.numel()
+        total_frames += current_frames
+        data_view = tensordict_data.reshape(-1)
+        replay_buffer.extend(data_view)
+
+        training_tds = []
+        training_start = time.time()
+        for _ in range(cfg.train.num_epochs):
+            for _ in range(cfg.collector.frames_per_batch // cfg.train.minibatch_size):
+                subdata = replay_buffer.sample()
+                loss_vals = loss_module(subdata)
+                training_tds.append(loss_vals.detach())
+
+                loss_value = loss_vals["loss"]
+
+                loss_value.backward()
+
+                total_norm = torch.nn.utils.clip_grad_norm_(
+                    loss_module.parameters(), cfg.train.max_grad_norm
+                )
+                training_tds[-1].set("grad_norm", total_norm.mean())
+
+                optim.step()
+                optim.zero_grad()
+                target_net_updater.step()
+
+        qnet_explore[1].step(frames=current_frames)  # Update exploration annealing
+        collector.update_policy_weights_()
+
+        training_time = time.time() - training_start
+
+        iteration_time = sampling_time + training_time
+        total_time += iteration_time
+        training_tds = torch.stack(training_tds)
+
+        # More logs
+        if cfg.logger.backend:
+            log_training(
+                logger,
+                training_tds,
+                tensordict_data,
+                sampling_time,
+                training_time,
+                total_time,
+                i,
+                current_frames,
+                total_frames,
+                step=i,
+            )
+
+        if (
+            cfg.eval.evaluation_episodes > 0
+            and i % cfg.eval.evaluation_interval == 0
+            and cfg.logger.backend
+        ):
+            evaluation_start = time.time()
+            with torch.no_grad(), set_exploration_type(ExplorationType.DETERMINISTIC):
+                env_test.frames = []
+                rollouts = env_test.rollout(
+                    max_steps=cfg.env.max_steps,
+                    policy=qnet,
+                    callback=rendering_callback,
+                    auto_cast_to_device=True,
+                    break_when_any_done=False,
+                    # We are running vectorized evaluation we do not want it to stop when just one env is done
+                )
+
+                evaluation_time = time.time() - evaluation_start
+
+                log_evaluation(logger, rollouts, env_test, evaluation_time, step=i)
+
+        if cfg.logger.backend == "wandb":
+            logger.experiment.log({}, commit=True)
+        sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
+
+
+if __name__ == "__main__":
+    train()