GridWorld (Lean-native finite RL environment)

This file defines a small deterministic GridWorld environment in TorchLean’s spec layer, along with two induced “MDP views”:

• a Spec.RL.Env view with explicit latent state,
• a Spec.RL.FiniteMDP view (deterministic finite-state discounted MDP),
• and a Spec.RL.FiniteStochastic.MDP view where transitions are represented as one-hot row-stochastic kernels.

This is intended as a Lean-native testbed for RL algorithm development and proofs: small enough to keep the transition proof small, but shaped like the objects used in standard RL theory.

References (high-level context only):

• Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed.), Chapter 3 (gridworld examples, discounted returns, Bellman operators).
• Puterman, Markov Decision Processes (1994), Chapters 6–7 (finite discounted MDPs).
• Gymnasium and TorchRL are useful API reference points for the environment/rollout shape: https://gymnasium.farama.org/ and https://pytorch.org/rl/

Environment Dynamics

Dynamics are deterministic and border-clamped:

• attempting to move outside the grid keeps the coordinate unchanged.

Reward / termination scheme:

• If the agent is already at the goal, the environment remains terminal (terminated = true) and yields reward 0.
• Otherwise, a step yields reward 0 iff the successor state is the goal, and reward -1 otherwise. The terminated flag matches “successor is goal”.
• truncated is always false (no time-limit semantics in this environment).

structure Spec.RL.Envs.GridWorld (width height : ℕ) : Type

A small deterministic GridWorld with a start cell, goal cell, and discount factor.

• start : Fin height × Fin width

  Initial state returned by reset.

• goal : Fin height × Fin width

  Terminal goal cell.

• discount : ℝ

  Discount factor $\gamma$ used by induced MDP views.

@[reducible, inline]

abbrev Spec.RL.Envs.GridWorld.State (width height : ℕ) : Type

A GridWorld state, represented by its row and column.

@[reducible, inline]

abbrev Spec.RL.Envs.GridWorld.Action : Type

A four-neighborhood GridWorld action.

def Spec.RL.Envs.GridWorld.Action.up : Action

Move one row upward.

def Spec.RL.Envs.GridWorld.Action.down : Action

Move one row downward.

def Spec.RL.Envs.GridWorld.Action.left : Action

Move one column to the left.

def Spec.RL.Envs.GridWorld.Action.right : Action

Move one column to the right.

def Spec.RL.Envs.GridWorld.rowUp {height : ℕ} (row : Fin height) : Fin height

The next row when moving one step up (saturating at 0).

def Spec.RL.Envs.GridWorld.rowDown {height : ℕ} (row : Fin height) : Fin height

The next row when moving one step down (clamped at height-1).

def Spec.RL.Envs.GridWorld.colLeft {width : ℕ} (col : Fin width) : Fin width

The next column when moving one step left (saturating at 0).

def Spec.RL.Envs.GridWorld.colRight {width : ℕ} (col : Fin width) : Fin width

The next column when moving one step right (clamped at width-1).

def Spec.RL.Envs.GridWorld.nextState {width height : ℕ} (state : State width height) (action : Action) : State width height

Deterministic successor state (border-clamped).

def Spec.RL.Envs.GridWorld.step {width height : ℕ} (gw : GridWorld width height) (state : State width height) (action : Action) : StepResult (State width height) ℝ

One deterministic step, returning reward and termination flags.

def Spec.RL.Envs.GridWorld.toEnv {width height : ℕ} (gw : GridWorld width height) : Env (State width height) Action (State width height) ℝ

Spec.RL.Env view of GridWorld with observations equal to latent states.

Finite-state MDP Views

To connect GridWorld to TorchLean’s finite discounted MDP layer we flatten the coordinate state:

Fin height × Fin width ≃ Fin (height * width).

This is the standard row-major encoding used throughout mathlib.

def Spec.RL.Envs.GridWorld.encode {width height : ℕ} (pos : State width height) : Fin (height * width)

Flatten a (row,col) grid coordinate into Fin (height * width).

def Spec.RL.Envs.GridWorld.decode {width height : ℕ} (i : Fin (height * width)) : State width height

Unflatten a Fin (height * width) state index into a (row,col) grid coordinate.

def Spec.RL.Envs.GridWorld.toFiniteMDP {width height : ℕ} (gw : GridWorld width height) : FiniteMDP ℝ (height * width) 4

Deterministic finite-state discounted MDP view of GridWorld.

def Spec.RL.Envs.GridWorld.oneHot {width height : ℕ} (next : Fin (height * width)) : TorchLean.Tensor ℝ [height * width]

One-hot transition kernel for a deterministic next state.

def Spec.RL.Envs.GridWorld.toFiniteStochasticMDP {width height : ℕ} (gw : GridWorld width height) : FiniteStochastic.MDP (height * width) 4

Finite-stochastic MDP view of GridWorld where $P(\mathord{\cdot}\mid s,a)$ is a one-hot row.