Finite Stochastic Discounted MDPs

This module extends TorchLean's finite deterministic MDP layer with finite-state stochastic transitions.

We stay in a deliberately small setting:

• finitely many states and actions,
• real-valued rewards and discount factor,
• row-stochastic transition kernels represented as typed vectors.

This is enough to formalize the Bellman expectation and optimality operators in the standard discounted setting without immediately introducing full measure-theoretic probability.

References:

• Bellman, Dynamic Programming (1957)
• Puterman, Markov Decision Processes (1994)
• Sutton and Barto, Reinforcement Learning: An Introduction
• TorchRL documentation for practical stochastic-environment / rollout APIs: https://pytorch.org/rl/

Naming note:

• The short names in this file live under Spec.RL.FiniteStochastic. Thus MDP, Valid, actionValue, and the Bellman operators mean the finite stochastic versions, not the deterministic tensor MDPs from Spec.RL.MDP or the Markov-kernel MDPs from Spec.RL.MarkovMDP.
• We use the file name FiniteStochasticMDP.lean to make the layer clear in imports, but keep the structure name MDP inside the namespace. The qualified name Spec.RL.FiniteStochastic.MDP is clearer at call sites than repeating the layer name twice.

structure Spec.RL.FiniteStochastic.MDP (nStates nActions : ℕ) : Type

Finite discounted MDP with stochastic next-state transitions.

• initialState : Fin nStates

  Canonical reset state.

• transitionProb : Fin nStates → Fin nActions → Tensor ℝ (Shape.dim nStates Shape.scalar)

  Transition probabilities P(. | s, a) over the finite next-state space.

• reward : Fin nStates → Fin nActions → ℝ

  Immediate reward r(s, a).

• terminated : Fin nStates → Fin nActions → Bool

  Task-defined terminal flag for (s, a).

• discount : ℝ

  Discount factor.

structure Spec.RL.FiniteStochastic.Valid {nStates nActions : ℕ} (mdp : MDP nStates nActions) : Prop

Well-formedness assumptions for a finite stochastic MDP.

• transition_nonneg (state : Fin nStates) (action : Fin nActions) (nextState : Fin nStates) : 0 ≤ (mdp.transitionProb state action).vecGet nextState

  Transition probabilities are nonnegative.

• transition_sums_to_one (state : Fin nStates) (action : Fin nActions) : ∑ nextState : Fin nStates, (mdp.transitionProb state action).vecGet nextState = 1

  Each transition row sums to 1.

• discount_nonneg : 0 ≤ mdp.discount

  Discount factor is nonnegative.

• discount_lt_one : mdp.discount < 1

  Discount factor is strictly less than 1.

def Spec.RL.FiniteStochastic.expectedNextValue {nStates nActions : ℕ} (mdp : MDP nStates nActions) (values : ValueFunction ℝ nStates) (state : Fin nStates) (action : Fin nActions) : ℝ

Expected next-state value under P(. | s, a).

def Spec.RL.FiniteStochastic.actionValue {nStates nActions : ℕ} (mdp : MDP nStates nActions) (values : ValueFunction ℝ nStates) (state : Fin nStates) (action : Fin nActions) : ℝ

Bellman-style state-action value induced by a candidate value function.

def Spec.RL.FiniteStochastic.actionValues {nStates nActions : ℕ} (mdp : MDP nStates nActions) (values : ValueFunction ℝ nStates) (state : Fin nStates) : Tensor ℝ (Shape.dim nActions Shape.scalar)

All state-action values Q_v(s, ·) for a fixed state and candidate value function.

def Spec.RL.FiniteStochastic.bellmanPolicy {nStates nActions : ℕ} (mdp : MDP nStates nActions) (policy : Policy nStates nActions) (values : ValueFunction ℝ nStates) : ValueFunction ℝ nStates

Bellman expectation operator for a deterministic policy.

def Spec.RL.FiniteStochastic.bellmanOptimality {nStates nActions : ℕ} [Fact (0 < nActions)] (mdp : MDP nStates nActions) (values : ValueFunction ℝ nStates) : ValueFunction ℝ nStates

Bellman optimality operator for a finite stochastic MDP.