Core Reinforcement-Learning Definitions

This module collects the small mathematical definitions that sit underneath TorchLean's RL development.

These definitions are intentionally spec-level rather than runtime-level:

• Bellman-style backups,
• discounted returns,
• generalized advantage estimation (GAE),
• and simple typed rollout records.

That keeps the actual RL mathematics in a proof-friendly namespace and avoids duplicating it inside runtime/trainer code.

Numerical Containers

Trajectories use Tensor α [horizon], including when the horizon is chosen at runtime. Rewards, values, and termination markers share that index, so mismatched trajectories cannot be silently truncated. Scalar recurrences and their evaluation order are explicit below.

Primary references:

• Sutton, "Learning to Predict by the Methods of Temporal Differences" (1988): https://doi.org/10.1023/A:1022633531479
• Watkins and Dayan, "Q-learning" (1992): https://doi.org/10.1007/BF00992698
• Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed.): http://incompleteideas.net/book/the-book-2nd.html
• Schulman et al., "High-Dimensional Continuous Control Using Generalized Advantage Estimation" (2015): https://arxiv.org/abs/1506.02438
• TorchRL documentation (rollouts, tensordicts, and GAE-style objectives): https://pytorch.org/rl/

def Spec.RL.continueMask {α : Type} [Zero α] [One α] (done : Bool) : α

Convert a terminal flag into a multiplicative continuation mask (1 for continue, 0 for stop).

def Spec.RL.discountedBackup {α : Type} [Zero α] [One α] [Add α] [Mul α] (reward gamma bootstrap : α) (done : Bool) : α

Bellman-style one-step backup: $r+\gamma(1-\mathtt{done})\mathtt{bootstrap}$.

def Spec.RL.tdTarget {α : Type} [Zero α] [One α] [Add α] [Mul α] (reward gamma nextValue : α) (done : Bool) : α

One-step TD target for state-value or action-value updates.

def Spec.RL.tdResidual {α : Type} [Zero α] [One α] [Add α] [Mul α] [Sub α] (value reward gamma nextValue : α) (done : Bool) : α

TD residual / Bellman error: $r+\gamma(1-d)\mathtt{nextValue}-\mathtt{value}$.

Shape-indexed trajectory calculations

def Spec.RL.discountedReturnsFrom {α : Type} [TorchLean.Storage α] [Zero α] [Add α] [Mul α] {n : ℕ} (gamma : α) (rewards : TorchLean.Tensor α [n]) (bootstrap : α := 0) : TorchLean.Tensor α [n]

Discounted returns with a far-right bootstrap, evaluated from right to left.

def Spec.RL.discountedReturns {α : Type} [TorchLean.Storage α] [Zero α] [Add α] [Mul α] {n : ℕ} (gamma : α) (rewards : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Discounted returns for a terminal trajectory.

def Spec.RL.discountedReturnsDone {α : Type} [TorchLean.Storage α] [Zero α] [One α] [Add α] [Mul α] {n : ℕ} (gamma : α) (rewards : TorchLean.Tensor α [n]) (dones : TorchLean.Tensor Bool [n]) (bootstrap : α := 0) : TorchLean.Tensor α [n]

Discounted returns with one termination marker per reward; unequal lengths are unrepresentable.

The multiplication order matches discountedBackup, including its floating-point behavior.

def Spec.RL.generalizedAdvantageEstimation {α : Type} [TorchLean.Storage α] [Zero α] [One α] [Add α] [Mul α] [Sub α] {n : ℕ} (gamma lam : α) (rewards values nextValues : TorchLean.Tensor α [n]) (dones : TorchLean.Tensor Bool [n]) : TorchLean.Tensor α [n]

Generalized Advantage Estimation, retaining the common horizon in all five tensor types.

def Spec.RL.returnsFromAdvantages {α : Type} [TorchLean.Storage α] [Add α] {n : ℕ} (advantages values : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Lambda-returns R_t = A_t + V_t, with equal lengths enforced by the tensor shape.