actors.py

import torch
import torch.nn as nn


class MLPActor(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, hidden_dim: int = 256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh(),
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        return self.net(obs)


class TierAwareActor(nn.Module):
    """
    CNN-based actor that processes tier-stratified observations.

    Supports both standard observations (K×M = num_tiers×4 dims) and
    age-augmented observations (K×M×2 = num_tiers×4×2 dims).
    """
    def __init__(self, num_tiers: int = 8, action_dim: int = 2, obs_dim: int = None):
        super().__init__()
        self.num_tiers = int(num_tiers)
        # Compute in_channels from obs_dim: 4 for standard, 8 for age-augmented
        if obs_dim is None:
            obs_dim = num_tiers * 4  # Default: standard observation
        self.obs_dim = obs_dim
        self.in_channels = obs_dim // num_tiers  # 4 for 32-dim, 8 for 64-dim

        self.net = nn.Sequential(
            nn.Conv1d(in_channels=self.in_channels, out_channels=16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(32 * self.num_tiers, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim),
            nn.Tanh(),
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        # Handle both batched and unbatched inputs
        if obs.dim() == 1:
            obs = obs.unsqueeze(0)
            squeeze_output = True
        else:
            squeeze_output = False
        # Reshape: (batch, obs_dim) → (batch, in_channels, num_tiers)
        x = obs.view(-1, self.num_tiers, self.in_channels).permute(0, 2, 1)
        out = self.net(x)
        if squeeze_output:
            out = out.squeeze(0)
        return out