Policy-Gradient Objectives

This module exposes typed helpers for the main categorical-policy objectives that modern policy-gradient code tends to rely on:

• REINFORCE,
• advantage actor-critic,
• trust-region / KL-penalized policy-gradient helpers,
• entropy regularization,
• soft actor-critic policy terms,
• PPO's clipped surrogate.

The helpers operate on logits for a finite action space and stay purely functional so they can be used from either eager runtime code or proof-oriented spec code.

Primary references:

• Williams, "Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning" (1992): https://doi.org/10.1023/A:1022672621406
• Mnih et al., "Asynchronous Methods for Deep Reinforcement Learning" (2016): https://arxiv.org/abs/1602.01783
• Schulman et al., "Trust Region Policy Optimization" (2015): https://arxiv.org/abs/1502.05477
• Schulman et al., "Proximal Policy Optimization Algorithms" (2017): https://arxiv.org/abs/1707.06347
• Schulman et al., "High-Dimensional Continuous Control Using Generalized Advantage Estimation" (2015): https://arxiv.org/abs/1506.02438

def Runtime.RL.PolicyGradient.actionPolicy {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) : TorchLean.Tensor α [nActions]

Softmax policy induced by a vector of logits.

def Runtime.RL.PolicyGradient.actionProbability {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (epsilon : α := Context.defaultEpsilon) : α

Probability of a selected action under a categorical policy.

def Runtime.RL.PolicyGradient.actionLogProbability {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (epsilon : α := Context.defaultEpsilon) : α

Log-probability of a selected action.

def Runtime.RL.PolicyGradient.entropyBonus {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (epsilon : α := Context.defaultEpsilon) : α

Guarded entropy bonus -Σ q(a) log q(a), where p = softmax logits and q = clamp p epsilon (1 - epsilon).

Clamped probabilities are used both as weights and as logarithm inputs, without renormalization. For positive epsilon, this can differ from the Shannon entropy of p.

def Runtime.RL.PolicyGradient.reinforceLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (returnOrAdvantage : α) (epsilon : α := Context.defaultEpsilon) : α

REINFORCE loss for one sampled action: -G_t * log π(a_t | s_t).

def Runtime.RL.PolicyGradient.actorLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (advantage : α) (epsilon : α := Context.defaultEpsilon) : α

Advantage actor-critic policy loss: -A_t * log π(a_t | s_t).

def Runtime.RL.PolicyGradient.criticLoss {α : Type} [Context α] (valuePrediction valueTarget : α) (valueCoef : α := 1) : α

Value-regression loss used by actor-critic and PPO critics.

def Runtime.RL.PolicyGradient.actorCriticLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (advantage valuePrediction valueTarget : α) (valueCoef : α := 1) (entropyCoef : α := 0) (epsilon : α := Context.defaultEpsilon) : α

Combined advantage actor-critic loss: policy term + value regression - entropy bonus.

def Runtime.RL.PolicyGradient.a2cLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (logits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (advantage valuePrediction valueTarget : α) (valueCoef : α := 1) (entropyCoef : α := 0) (epsilon : α := Context.defaultEpsilon) : α

Advantage actor-critic loss with an explicit entropy bonus coefficient.

This is the A2C/A3C-shaped single-sample objective: -A_t log π(a_t|s_t) + c_v value_loss - c_e H(π(.|s_t)).

def Runtime.RL.PolicyGradient.importanceRatio {α : Type} [Context α] (newLogProb oldLogProb : α) : α

Importance ratio π_new(a|s) / π_old(a|s) computed from log-probabilities.

def Runtime.RL.PolicyGradient.categoricalKL {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (oldProbs newProbs : TorchLean.Tensor α [nActions]) (epsilon : α := Context.defaultEpsilon) : α

Categorical KL divergence after clamping and normalizing both probability vectors.

For finite inputs and 0 < epsilon < 1/2, clamp each vector into [epsilon, 1-epsilon], then divide by its sum before computing Σ_a old(a) * (log old(a) - log new(a)). Normalization keeps the guarded inputs probability distributions. The result agrees with KL(old || new) when clamping leaves normalized inputs unchanged, up to floating-point rounding. Empty action vectors return zero.

def Runtime.RL.PolicyGradient.categoricalKLFromLogits {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (oldLogits newLogits : TorchLean.Tensor α [nActions]) (epsilon : α := Context.defaultEpsilon) : α

Categorical KL divergence from logits, using the clamped, normalized policies of categoricalKL.

def Runtime.RL.PolicyGradient.trpoSurrogateFromRatio {α : Type} [Context α] (ratio advantage : α) : α

TRPO-style surrogate objective from a precomputed importance ratio: ratio * A.

TRPO maximizes this surrogate subject to a KL trust-region constraint. We expose the scalar surrogate separately from the constraint so callers can choose line search / penalty / diagnostics.

def Runtime.RL.PolicyGradient.klPenalizedPolicyLoss {α : Type} [Context α] (ratio advantage kl penaltyCoef : α) : α

KL-penalized policy-gradient loss: -(ratio * A) + β * KL(old || new).

This is not the full constrained TRPO optimizer; it is the differentiable scalar objective commonly used as a practical surrogate or diagnostic when implementing trust-region updates.

def Runtime.RL.PolicyGradient.sacCategoricalActorLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} [NeZero nActions] (logits qValues : TorchLean.Tensor α [nActions]) (temperature : α) : α

Finite-action SAC actor objective, minimized over actor logits: ∑ a, π(a|s) * (temperature * log π(a|s) - Q(s,a)).

Max-shifted exponentials supply normalized policy weights. When a negative logit and positive maximum could overflow their difference, weight each before subtracting. Otherwise keep the usual shifted expression. Temperature multiplies the weighted entropy term. These rearrangements retain differentiation through the weights and normalization without taking the log of a zero probability.

For an actor-only gradient, callers hold qValues and temperature constant with respect to actor parameters. If two critics are used, pass their pointwise minimum as qValues. This function sums over actions for one state without averaging across states.

def Runtime.RL.PolicyGradient.ppoClippedObjectiveFromRatio {α : Type} [Context α] (ratio advantage clipEps : α) : α

PPO clipped surrogate objective from a precomputed importance ratio:

min(ratio * A, clip(ratio, 1-ε, 1+ε) * A).

This helper is useful when you already have the ratio (e.g. from cached log-probabilities) and want to avoid recomputing it from logits.

def Runtime.RL.PolicyGradient.ppoClippedObjective {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (newLogits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (oldLogProb advantage clipEps : α) (epsilon : α := Context.defaultEpsilon) : α

PPO clipped surrogate objective for one sampled action.

This is the objective to maximize: min(r_t A_t, clip(r_t, 1-ε, 1+ε) A_t).

def Runtime.RL.PolicyGradient.ppoLoss {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} (newLogits : TorchLean.Tensor α [nActions]) (action : Fin nActions) (oldLogProb advantage valuePrediction valueTarget clipEps : α) (valueCoef : α := 1) (entropyCoef : α := 0) (epsilon : α := Context.defaultEpsilon) : α

PPO loss to minimize: -L_clip + c_v * value_loss - c_e * entropy.

def Runtime.RL.PolicyGradient.sampleCategorical {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} [NeZero nActions] (seed counter : ℕ) (probs : TorchLean.Tensor α [nActions]) : ℕ × Fin nActions

Sample from a categorical distribution represented as a probability vector.

seed and counter form an explicit RNG stream identifier. The function returns the incremented counter together with the sampled action index.

Implementation note: this uses the standard cumulative-sum / inverse-CDF sampler.

def Runtime.RL.PolicyGradient.sampleActionFromLogits {α : Type} [TorchLean.Storage α] [Context α] {nActions : ℕ} [NeZero nActions] (seed counter : ℕ) (logits : TorchLean.Tensor α [nActions]) : ℕ × Fin nActions

Sample an action from logits by applying softmax then sampleCategorical.