QMIX_PS.py

"""
Single-file implementation of QMIX-PS.

QMIX-PS: QMIX with Parameter Sharing.
This module implements an off-policy value-based multi-agent learner with:
- a shared local Q-network per agent observation, and
- a monotonic mixing network that maps per-agent Q-values to a global Q_tot.
"""

import random
import numpy as np
import torch
import torch.nn as nn
import gymnasium as gym

from imprl.agents.primitives.exploration_schedulers import LinearExplorationScheduler
from imprl.agents.primitives.replay_memory import AbstractReplayMemory
from imprl.agents.primitives.MLP import NeuralNetwork


class QMixer(nn.Module):
    def __init__(self, n_agents, state_dim, mixer_config):
        super().__init__()

        """
        Mixing network for QMIX

        mixer: Q_tot = f(q_values @ w1 + b1) @ w2 + b2

        w1, b1, w2, b2 are mixing parameters generated by hypernetworks
        using the state as input. All mixing parameters are strictly positive
        to ensure monotonicity of the mixing function.
        """

        self.n_agents = n_agents
        self.hypernet_hidden_units = mixer_config["hypernet_hidden_units"]
        self.mixer_embed_dim = mixer_config["mixer_embed_dim"]

        ## Hypernetworks
        # hypernetworks for w1 and w2, depending on the number of hidden units
        if self.hypernet_hidden_units in [0, None]:
            self.hypernet_w1 = nn.Linear(state_dim, self.mixer_embed_dim * n_agents)
            self.hypernet_w2 = nn.Linear(state_dim, self.mixer_embed_dim)
        elif self.hypernet_hidden_units > 0:
            self.hypernet_w1 = nn.Sequential(
                nn.Linear(state_dim, self.hypernet_hidden_units),
                nn.ReLU(),
                nn.Linear(
                    self.hypernet_hidden_units, self.mixer_embed_dim * n_agents
                ),
            )
            self.hypernet_w2 = nn.Sequential(
                nn.Linear(state_dim, self.hypernet_hidden_units),
                nn.ReLU(),
                nn.Linear(self.hypernet_hidden_units, self.mixer_embed_dim),
            )
        else:
            raise ValueError("Invalid number of hidden units")

        self.hypernet_b1 = nn.Linear(state_dim, self.mixer_embed_dim)
        self.hypernet_b2 = nn.Sequential(
            nn.Linear(state_dim, self.mixer_embed_dim),
            nn.ReLU(),
            nn.Linear(self.mixer_embed_dim, 1),
        )

    def forward(self, q_values, states):
        q_values = q_values.view(-1, 1, self.n_agents)

        w1 = torch.abs(self.hypernet_w1(states))
        w1 = w1.view(-1, self.n_agents, self.mixer_embed_dim)

        b1 = self.hypernet_b1(states)
        b1 = b1.view(-1, 1, self.mixer_embed_dim)

        w2 = torch.abs(self.hypernet_w2(states))
        w2 = w2.view(-1, self.mixer_embed_dim, 1)

        b2 = self.hypernet_b2(states)
        b2 = b2.view(-1, 1, 1)

        # First layer
        # a = f(q @ w1 + b1)
        z = torch.bmm(q_values, w1) + b1
        a = torch.nn.functional.elu(z)

        # Second layer
        # q_tot = a @ w2 + b2
        q_total = torch.bmm(a, w2) + b2
        return q_total.view(-1, 1)


class QMIXParameterSharing:
    name = "QMIX-PS"  # display names used by experiment scripts/loggers.
    full_name = "QMIX with Parameter Sharing"

    # Algorithm taxonomy.
    paradigm = "CTDE"
    formulation = "Dec-POMDP"
    algorithm_type = "value-based"
    policy_regime = "off-policy"
    policy_type = "epsilon-greedy"

    # Training/runtime properties.
    uses_replay_memory = True
    parameter_sharing = True
    collect_state_info = True

    def __init__(self, env, config, device):
        """Initialize shared Q/mixer modules, replay, exploration, and counters."""
        assert (
            env.__class__.__name__ == "MultiAgentWrapper"
        ), "QMIX-PS only supports multi-agent environments"

        # ---------- Core references and counters ----------
        self.env, self.device, self.cfg = env, device, config
        self.episode = 0
        self.total_time = 0
        self.time = 0
        self.episode_return = 0

        # ---------- Evaluation discount + replay ----------
        try:
            self.eval_discount_factor = env.core.discount_factor
        except AttributeError:
            self.eval_discount_factor = 1.0

        self.replay_memory = AbstractReplayMemory(self.cfg["MAX_MEMORY_SIZE"])
        # Set a short warmup period to populate replay before training starts.
        self.warmup_threshold = 10 * self.cfg["BATCH_SIZE"]

        # ---------- Exploration ----------
        self.exploration_strategy = self.cfg["EXPLORATION_STRATEGY"]
        self.exploration_param = self.exploration_strategy["max_value"]
        self.exp_scheduler = LinearExplorationScheduler(
            self.exploration_strategy["min_value"],
            num_episodes=self.exploration_strategy["num_episodes"],
        )

        # ---------- QMIX config ----------
        self.use_state_info = self.cfg["USE_STATE_INFO"]
        self.mixer_config = self.cfg["MIXER_CONFIG"]
        self.target_network_reset = self.cfg["TARGET_NETWORK_RESET"]

        # ---------- Environment dimensions ----------
        self.n_agent_actions = env.action_space.n
        self.n_agents = env.n_agents

        obs, _ = env.reset()
        local_obs = env.multiagent_idx_percept(obs)
        n_inputs = local_obs.shape[-1]
        state_dim = gym.spaces.utils.flatdim(env.perception_space)

        # ---------- Network architectures ----------
        network_arch = [n_inputs] + self.cfg["NETWORK_CONFIG"]["hidden_layers"]
        network_arch += [self.n_agent_actions]

        # ---------- Q / mixer modules ----------
        self.q_network = NeuralNetwork(network_arch, initialization="orthogonal").to(
            device
        )
        self.q_mixer = QMixer(self.n_agents, state_dim, self.mixer_config).to(device)

        self.target_network = NeuralNetwork(network_arch).to(device)
        self.target_mixer = QMixer(self.n_agents, state_dim, self.mixer_config).to(
            device
        )
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.target_mixer.load_state_dict(self.q_mixer.state_dict())

        # One optimizer/scheduler for both local-Q and mixer parameters.
        all_params = list(self.q_network.parameters()) + list(self.q_mixer.parameters())
        self.optimizer = getattr(torch.optim, self.cfg["NETWORK_CONFIG"]["optimizer"])(
            all_params, lr=self.cfg["NETWORK_CONFIG"]["lr"]
        )
        lrs = self.cfg["NETWORK_CONFIG"]["lr_scheduler"]
        self.lr_scheduler = getattr(torch.optim.lr_scheduler, lrs["scheduler"])(
            self.optimizer, **lrs["kwargs"]
        )

        # ---------- Logging fields ----------
        self.logger = {
            "TD_loss": None,
            "learning_rate": self.cfg["NETWORK_CONFIG"]["lr"],
            "exploration_param": self.exploration_param,
        }

    def reset_episode(self, training=True):
        """Reset episode counters; update epsilon/LR and sync target networks on schedule."""
        self.episode_return = 0
        self.episode += 1
        self.time = 0

        if not training:
            return

        # Epsilon schedule for behavior policy.
        self.exploration_param = self.exp_scheduler.step()
        self.logger["exploration_param"] = self.exploration_param

        # Start scheduler/target updates only when replay batches are trainable.
        if self.total_time > self.warmup_threshold:
            if self.episode % self.target_network_reset == 0:
                self.target_network.load_state_dict(self.q_network.state_dict())
                self.target_mixer.load_state_dict(self.q_mixer.state_dict())
            self.lr_scheduler.step()
            self.logger["learning_rate"] = self.lr_scheduler.get_last_lr()[0]

    def select_action(self, observation, training):
        # Epsilon-greedy exploration during training.
        if (
            training
            and self.exploration_strategy["name"] == "epsilon_greedy"
            and self.exploration_param > random.random()
        ):
            action = [self.env.action_space.sample() for _ in range(self.n_agents)]
            return action, torch.tensor(action)

        # Greedy local actions from shared Q-network.
        local_obs = self.env.multiagent_idx_percept(observation)
        t_local_obs = torch.as_tensor(local_obs, device=self.device)
        q_values = self.q_network.forward(t_local_obs, training=training)
        t_action = torch.argmax(q_values, dim=-1)
        action = t_action.cpu().detach().numpy()
        if training:
            return action, t_action.detach().cpu()
        return action

    def process_experience(
        self,
        observation,
        state,
        t_action,
        next_observation,
        next_state,
        reward,
        terminated,
        truncated,
    ):
        """Store transition, run one replay update after warmup, and log episode end."""
        # Update episode return/time counters.
        self.process_rewards(reward)

        # Replay stores both local and centralized representations for QMIX.
        self.replay_memory.store_experience(
            self.env.multiagent_idx_percept(observation),
            self.env.system_percept(observation),
            self.env.system_percept(state),
            t_action,
            self.env.multiagent_idx_percept(next_observation),
            self.env.system_percept(next_observation),
            self.env.system_percept(next_state),
            reward,
            terminated,
            truncated,
        )

        # Train from replay after a short warmup.
        if self.total_time > self.warmup_threshold:
            sample_batch = self.replay_memory.sample_batch(self.cfg["BATCH_SIZE"])
            self.train(*sample_batch)

        # Episode-level logging.
        if terminated or truncated:
            self.logger["episode"] = self.episode
            self.logger["episode_cost"] = -self.episode_return

    def compute_loss(self, *args):
        """
        Compute TD loss for one replay batch using QMIX mixing.

        Core steps:
        - Current Q_tot:
          gather local Q_i(o_i, a_i), then mix with QMixer.
        - Future Q_tot:
          pick greedy per-agent actions from target local-Q outputs,
          gather target Q_i(o'_i, a'_i), then mix with target mixer.
        - TD target:
          y = r + gamma * (1 - done) * Q_tot_target(s', a')
        """
        (
            ma_beliefs,
            beliefs,
            states,
            actions,
            ma_next_beliefs,
            next_beliefs,
            next_states,
            rewards,
            terminations,
            _truncations,
        ) = args

        # ---------- Tensorize replay samples ----------
        t_ma_beliefs = torch.as_tensor(np.asarray(ma_beliefs), device=self.device)
        t_beliefs = torch.as_tensor(np.asarray(beliefs), device=self.device)
        t_states = torch.as_tensor(np.asarray(states), device=self.device)
        t_actions = torch.stack(actions).to(self.device)
        t_ma_next_beliefs = torch.as_tensor(
            np.asarray(ma_next_beliefs), device=self.device
        )
        t_next_beliefs = torch.as_tensor(np.asarray(next_beliefs), device=self.device)
        t_next_states = torch.as_tensor(np.asarray(next_states), device=self.device)
        t_rewards = torch.as_tensor(rewards, device=self.device).reshape(-1, 1)
        t_terminations = torch.as_tensor(
            terminations, dtype=torch.int, device=self.device
        ).reshape(-1, 1)

        # ---------- Current Q_tot estimates ----------
        all_q_values = self.q_network.forward(t_ma_beliefs)
        q_values = torch.gather(all_q_values, dim=2, index=t_actions.unsqueeze(2))
        q_values = q_values.squeeze(-1)
        mix_input = t_states if self.use_state_info else t_beliefs
        current_values = self.q_mixer.forward(q_values, mix_input)

        # ---------- Target Q_tot ----------
        with torch.no_grad():
            target_q_values = self.target_network.forward(t_ma_next_beliefs)
            best_actions = torch.argmax(target_q_values, dim=2)
            future_values = torch.gather(
                target_q_values, dim=2, index=best_actions.unsqueeze(2)
            ).squeeze(-1)
            target_mix_input = t_next_states if self.use_state_info else t_next_beliefs
            q_total_future = self.target_mixer.forward(future_values, target_mix_input)
            not_terminals = 1 - t_terminations
            td_targets = (
                t_rewards + self.cfg["DISCOUNT_FACTOR"] * q_total_future * not_terminals
            )

        return torch.nn.functional.mse_loss(current_values, td_targets)

    def train(self, *args):
        loss = self.compute_loss(*args)

        # Zero gradients, perform a backward pass, and update the weights.
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # logging value update
        self.logger["TD_loss"] = loss.detach()

    def process_rewards(self, reward):
        """Accumulate discounted episode return and advance time counters."""
        self.episode_return += reward * self.eval_discount_factor**self.time
        self.time += 1
        self.total_time += 1

    def report(self, stats=None):
        """Print episode-level progress to stdout."""
        print(f"Ep:{self.episode:05}| Cost: {-self.episode_return:.2f}", flush=True)
        if stats is not None:
            print(stats)

    def save_weights(self, path, episode):
        """Save local-Q and mixer parameters for a checkpoint id."""
        torch.save(self.q_network.state_dict(), f"{path}/q_network_{episode}.pth")
        torch.save(self.q_mixer.state_dict(), f"{path}/q_mixer_{episode}.pth")

    def load_weights(self, path, episode):
        """Load local-Q and mixer parameters from a checkpoint id."""
        self.q_network.load_state_dict(
            torch.load(
                f"{path}/q_network_{episode}.pth", map_location=torch.device("cpu")
            )
        )
        self.q_mixer.load_state_dict(
            torch.load(f"{path}/q_mixer_{episode}.pth", map_location=torch.device("cpu"))
        )