designing-policy-actor

name	Designing Policy Actor
description	This skill should be used when the user asks to "design a policy network", "implement an actor network", "build the action distribution layer", "implement Categorical distribution for RL", "implement Gaussian policy", "handle continuous action spaces", "implement multi-discrete actions", or "build the actor head for PPO/SAC/TD3". Do NOT apply Softmax directly to the network output before creating a Categorical distribution.
version	0.1.0

Designing Policy Actor

The Policy Actor maps latent state representations to action probability distributions. Correct distribution wiring prevents numerical collapse and ensures valid log-probability computations for gradient updates.

Discrete Action Spaces (Categorical)

The final layer yields N raw logits (one per discrete action). Pass them directly to torch.distributions.Categorical:

import torch
import torch.nn as nn
from torch.distributions import Categorical

class DiscreteActor(nn.Module):
    def __init__(self, features_dim: int, n_actions: int):
        super().__init__()
        self.head = nn.Sequential(
            nn.Linear(features_dim, 64),
            nn.Tanh(),
            nn.Linear(64, n_actions),  # Raw logits — no activation
        )

    def act(self, features, deterministic=False):
        logits = self.head(features)
        dist = Categorical(logits=logits)  # Never pass probs= after Softmax
        if deterministic:
            action = logits.argmax(dim=-1)
        else:
            action = dist.sample()
        log_prob = dist.log_prob(action)
        return action, log_prob

    def evaluate(self, features, action):
        logits = self.head(features)
        dist = Categorical(logits=logits)
        log_prob = dist.log_prob(action)
        entropy = dist.entropy()
        return log_prob, entropy

Never apply Softmax before instantiation. Categorical(logits=logits) applies log-softmax internally. Passing Categorical(probs=softmax(logits)) causes double-softmax → numerical collapse of the log-probabilities.

Continuous Action Spaces (Gaussian/Normal)

The final layer outputs μ (mean) and log_std. Use the script for the std bounding logic:

• scripts/gaussian_actor.py — Full Gaussian actor with proper log_std clamping

Key rules:

• Parameterize the network to output log_std (unbounded range) rather than std directly
• Convert: std = exp(log_std) or std = softplus(log_std) to ensure strict positivity
• Clamp log_std to reasonable bounds (e.g., [-20, 2]) to prevent degenerate distributions

from torch.distributions import Normal
import torch

class ContinuousActor(nn.Module):
    LOG_STD_MIN = -20
    LOG_STD_MAX = 2

    def __init__(self, features_dim: int, action_dim: int):
        super().__init__()
        self.mu_head = nn.Linear(features_dim, action_dim)
        self.log_std_head = nn.Linear(features_dim, action_dim)

    def act(self, features, deterministic=False):
        mu = self.mu_head(features)
        log_std = self.log_std_head(features).clamp(self.LOG_STD_MIN, self.LOG_STD_MAX)
        std = log_std.exp()
        dist = Normal(loc=mu, scale=std)
        if deterministic:
            action = mu
        else:
            action = dist.rsample()  # rsample() for reparameterization gradient
        log_prob = dist.log_prob(action).sum(dim=-1)
        return action, log_prob

    def evaluate(self, features, action):
        mu = self.mu_head(features)
        log_std = self.log_std_head(features).clamp(self.LOG_STD_MIN, self.LOG_STD_MAX)
        std = log_std.exp()
        dist = Normal(loc=mu, scale=std)
        log_prob = dist.log_prob(action).sum(dim=-1)
        entropy = dist.entropy().sum(dim=-1)
        return log_prob, entropy

Squashed Gaussian (SAC / Bounded Actions)

For environments with bounded action spaces (e.g., [-1, 1]), apply tanh squashing after sampling and correct the log-probability for the Jacobian:

action_raw = dist.rsample()
action = torch.tanh(action_raw)
# Correct log_prob for tanh squashing (Jacobian correction)
log_prob = dist.log_prob(action_raw).sum(-1) - torch.log(1 - action.pow(2) + 1e-6).sum(-1)

Multi-Discrete Action Spaces

For multiple independent discrete sub-actions (e.g., camera pan + zoom independently):

from torch.distributions import Categorical

class MultiDiscreteActor(nn.Module):
    def __init__(self, features_dim: int, nvec: list):
        """nvec: list of action counts per sub-action, e.g. [4, 3, 2]"""
        super().__init__()
        self.heads = nn.ModuleList([
            nn.Linear(features_dim, n) for n in nvec
        ])

    def act(self, features, deterministic=False):
        actions, log_probs = [], []
        for head in self.heads:
            logits = head(features)
            dist = Categorical(logits=logits)
            a = logits.argmax(-1) if deterministic else dist.sample()
            actions.append(a)
            log_probs.append(dist.log_prob(a))
        return torch.stack(actions, dim=-1), torch.stack(log_probs, dim=-1).sum(-1)

Required Interface

Every Actor must expose these two explicit functions:

Method	Signature	Purpose
act(features, deterministic)	Returns (action, log_prob)	Used during rollout collection
evaluate(features, action)	Returns (log_prob, entropy)	Used during gradient update step

Additional Resources

Scripts

• scripts/gaussian_actor.py — Full Gaussian actor with log_std bounding, softplus option, SAC squashing

References

• references/torch_distributions_api.md — torch.distributions API reference, distribution parameter shapes, common pitfalls