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}$.

Termination masks by multiplication, without short-circuiting the bootstrap. In floating-point arithmetic, a terminal transition can therefore still produce NaN from a NaN or infinite bootstrap.

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.discountedReturns {α : Type} [TorchLean.Storage α] [Zero α] [Add α] [Mul α] {n : ℕ} (gamma : α) (rewards : TorchLean.Tensor α [n]) (bootstrap : α := 0) : TorchLean.Tensor α [n]

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

The default bootstrap is 0, representing 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. Terminal branches do not multiply unused bootstrap values or future advantages by zero.

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.