1{
2  "op": "replace",
3  "filename": "epymarl/src/modules/critics/custom_critic.py",
4  "start_line": 7,
5  "end_line": 8,
6  "content": "# ── Custom imports (editable) ────────────────────────────────────────────\n"
7}

1{
2  "op": "replace",
3  "filename": "epymarl/src/modules/critics/custom_critic.py",
4  "start_line": 12,
5  "end_line": 68,
6  "content": "class CustomCritic(nn.Module):\n    \"\"\"Observation-Gated Centralized Critic (OGC²) for MAPPO.\n\n    Algorithmic contribution: Feature-wise Linear Modulation (FiLM) fusion\n    of global state and per-agent observations.\n\n    In cooperative MARL, each agent's value depends on both the global team\n    state AND the agent's local situation. The same global state has different\n    value implications for different agents depending on their local context\n    (e.g., a low-health unit should value retreat more than a full-health one).\n\n    Standard MAPPO ignores per-agent obs in the critic (state + agent_id only).\n    Simple concatenation of state + obs treats them symmetrically. Instead, we\n    use FiLM conditioning: each agent's obs encoder produces per-feature scale\n    (gamma) and shift (beta) that multiplicatively modulate the state\n    representation. This creates agent-specific \"lenses\" through which to\n    interpret the global state, capturing the interaction between global\n    context and local situation.\n\n    Architecture:\n      state  → Linear → ReLU → h_s (hidden state features)\n      obs+id → Linear → ReLU → h_o (hidden obs features)\n      h_o → Linear → Sigmoid → γ  (per-feature gate)\n      h_o → Linear           → β  (per-feature shift)\n      modulated = γ ⊙ h_s + β + h_o  (FiLM + residual)\n      LayerNorm(modulated) → Linear → ReLU → Linear → V\n    \"\"\"\n\n    def __init__(self, scheme, args):\n        super(CustomCritic, self).__init__()\n        self.args = args\n        self.n_agents = args.n_agents\n        self.n_actions = args.n_actions\n        self.output_type = \"v\"\n\n        obs_dim = int(scheme[\"obs\"][\"vshape\"])\n        state_dim = int(scheme[\"state\"][\"vshape\"])\n        H = args.hidden_dim\n\n        # Dual-stream first-layer encoders\n        self.state_fc = nn.Linear(state_dim, H)\n        self.obs_fc = nn.Linear(obs_dim + self.n_agents, H)\n\n        # FiLM generators: obs features → per-feature scale & shift for state\n        self.film_gamma = nn.Linear(H, H)   # → sigmoid → multiplicative gate\n        self.film_beta = nn.Linear(H, H)    # → additive shift\n\n        # Layer norm after modulation for training stability\n        self.ln = nn.LayerNorm(H)\n\n        # Value head (shared across agents)\n        self.fc2 = nn.Linear(H, H)\n        self.fc3 = nn.Linear(H, 1)\n\n    def forward(self, batch, t=None):\n        bs = batch.batch_size\n        max_t = batch.max_seq_length if t is None else 1\n        ts = slice(None) if t is None else slice(t, t + 1)\n\n        # Extract inputs\n        obs = batch[\"obs\"][:, ts]                                        # (B, T, n, obs_dim)\n        state = batch[\"state\"][:, ts]                                    # (B, T, state_dim)\n        state = state.unsqueeze(2).expand(-1, -1, self.n_agents, -1)     # (B, T, n, state_dim)\n\n        agent_id = th.eye(self.n_agents, device=batch.device)\n        agent_id = agent_id.unsqueeze(0).unsqueeze(0).expand(bs, max_t, -1, -1)\n\n        # Encode state stream\n        h_s = F.relu(self.state_fc(state))                               # (B, T, n, H)\n\n        # Encode obs stream (obs + agent one-hot)\n        obs_input = th.cat([obs, agent_id], dim=-1)                      # (B, T, n, obs+n)\n        h_o = F.relu(self.obs_fc(obs_input))                             # (B, T, n, H)\n\n        # FiLM: obs-conditioned gating of state features\n        gamma = th.sigmoid(self.film_gamma(h_o))                         # (B, T, n, H) ∈ (0,1)\n        beta = self.film_beta(h_o)                                       # (B, T, n, H)\n\n        # Modulate state features + residual from obs stream\n        modulated = gamma * h_s + beta + h_o                             # (B, T, n, H)\n        modulated = self.ln(modulated)\n\n        # Value head\n        x = F.relu(self.fc2(modulated))                                  # (B, T, n, H)\n        q = self.fc3(x)                                                  # (B, T, n, 1)\n        return q"
7}

1{
2  "content": "class CustomCritic(nn.Module):\n    \"\"\"Observation-Gated Centralized Critic (OGC²) for MAPPO.\n\n    Algorithmic contribution: Dual-path critic with observation-conditioned\n    gating of global state features, plus a stabilizing residual path.\n\n    In cooperative MARL, each agent's value depends on both the global team\n    state AND the agent's local situation. Standard MAPPO uses only\n    state + agent_id, ignoring per-agent observations in the critic.\n\n    Our approach: the agent's local observation generates a soft gate that\n    selectively amplifies or suppresses features of the global state encoding.\n    This lets each agent extract value-relevant information from the shared\n    state based on its own situation (e.g., a low-health agent values\n    retreat states differently than a full-health one).\n\n    Two paths are combined:\n      1. Gated path:  obs → gate → element-wise modulate state features\n      2. Direct path:  state + agent_id → standard encoding (stabilizer)\n    The combination provides both novelty (gating) and stability (direct path).\n    \"\"\"\n\n    def __init__(self, scheme, args):\n        super(CustomCritic, self).__init__()\n        self.args = args\n        self.n_agents = args.n_agents\n        self.n_actions = args.n_actions\n        self.output_type = \"v\"\n\n        obs_dim = int(scheme[\"obs\"][\"vshape\"])\n        state_dim = int(scheme[\"state\"][\"vshape\"])\n        H = args.hidden_dim\n\n        # ── Direct path: standard MAPPO-style (state + agent_id) ──\n        self.direct_fc1 = nn.Linear(state_dim + self.n_agents, H)\n\n        # ── Gated path: obs-conditioned modulation of state ──\n        # State encoder\n        self.state_fc = nn.Linear(state_dim, H)\n        # Obs encoder (generates gate)\n        self.obs_fc = nn.Linear(obs_dim + self.n_agents, H)\n        # Gate projection\n        self.gate_fc = nn.Linear(H, H)\n\n        # ── Fusion: combine both paths ──\n        self.fc2 = nn.Linear(H, H)\n        self.fc3 = nn.Linear(H, 1)\n\n    def forward(self, batch, t=None):\n        bs = batch.batch_size\n        max_t = batch.max_seq_length if t is None else 1\n        ts = slice(None) if t is None else slice(t, t + 1)\n\n        # Extract inputs\n        obs = batch[\"obs\"][:, ts]                                        # (B, T, n, obs_dim)\n        state = batch[\"state\"][:, ts]                                    # (B, T, state_dim)\n        state = state.unsqueeze(2).expand(-1, -1, self.n_agents, -1)     # (B, T, n, state_dim)\n\n        agent_id = th.eye(self.n_agents, device=batch.device)\n        agent_id = agent_id.unsqueeze(0).unsqueeze(0).expand(bs, max_t, -1, -1)\n\n        # ── Direct path (stabilizer): standard state + agent_id ──\n        direct_input = th.cat([state, agent_id], dim=-1)                 # (B, T, n, state+n)\n        h_direct = F.relu(self.direct_fc1(direct_input))                 # (B, T, n, H)\n\n        # ── Gated path: obs-conditioned state modulation ──\n        h_s = F.relu(self.state_fc(state))                               # (B, T, n, H)\n        obs_input = th.cat([obs, agent_id], dim=-1)                      # (B, T, n, obs+n)\n        h_o = F.relu(self.obs_fc(obs_input))                             # (B, T, n, H)\n        gate = th.sigmoid(self.gate_fc(h_o))                             # (B, T, n, H) ∈ (0,1)\n        h_gated = gate * h_s                                             # (B, T, n, H)\n\n        # ── Combine: gated features + direct path ──\n        combined = h_direct + h_gated                                    # (B, T, n, H)\n\n        # ── Value head ──\n        x = F.relu(self.fc2(combined))                                   # (B, T, n, H)\n        q = self.fc3(x)                                                  # (B, T, n, 1)\n        return q",
3  "end_line": 97,
4  "filename": "epymarl/src/modules/critics/custom_critic.py",
5  "op": "replace",
6  "start_line": 12
7}

1{
2  "n": 2
3}