{-# LANGUAGE DerivingStrategies #-}

-- | A tiny gridworld for the circuits-rl frontier spike.
--
-- The demonstration pins the reward/policy design choice upfront:
--
-- * Reward is a function of state, applied via a local 'scoreBy' modality.
-- * Transitions are deterministic for exact hand-checking.
-- * Value iteration is shown both directly and as composition in 'Prob'.
module Circuit.RL.GridWorld
  ( -- * Gridworld
    State (..),
    Action (..),
    step,
    reward,

    -- * Direct value iteration
    bellmanPolicy,
    bellmanOpt,
    valueIter,
    optimalPolicy,

    -- * Prob-composition view
    scoreBy,
    transP,
    rewardP,
    bellmanP,
    backupP,

    -- * Discounted-return oracle
    discountedReturn,
    closedFormReturn,

    -- * System (Prob) view
    expectSystem,
    gridSystem,
    mdpSystem,
    mdpCheck,
    pomdpSystem,
    pomdpCheck,
    Observation (..),
    observe,
    bellmanSystem,
    valueIterSystem,

    -- * Tropical / shortest-path row
    Tropical (..),
    shortestPath,
  )
where

import Circuit.Category (id, (.))
import Circuit.Poly (Mono, Poly (..))
import Circuit.Prob (Prob (..), embed, score)
import Circuit.System (System, monoDir, monoIn, runSystem, system)
import Data.List (foldl', maximumBy)
import Data.Ord (comparing)
import Data.Void (Void, absurd)
import Prelude hiding (id, (.))

-- | A one-dimensional chain of four states; 'Goal' is the absorbing target.
data State = S0 | S1 | S2 | Goal
  deriving stock (State -> State -> Bool
(State -> State -> Bool) -> (State -> State -> Bool) -> Eq State
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: State -> State -> Bool
== :: State -> State -> Bool
$c/= :: State -> State -> Bool
/= :: State -> State -> Bool
Eq, Int -> State -> ShowS
[State] -> ShowS
State -> String
(Int -> State -> ShowS)
-> (State -> String) -> ([State] -> ShowS) -> Show State
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> State -> ShowS
showsPrec :: Int -> State -> ShowS
$cshow :: State -> String
show :: State -> String
$cshowList :: [State] -> ShowS
showList :: [State] -> ShowS
Show, Int -> State
State -> Int
State -> [State]
State -> State
State -> State -> [State]
State -> State -> State -> [State]
(State -> State)
-> (State -> State)
-> (Int -> State)
-> (State -> Int)
-> (State -> [State])
-> (State -> State -> [State])
-> (State -> State -> [State])
-> (State -> State -> State -> [State])
-> Enum State
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: State -> State
succ :: State -> State
$cpred :: State -> State
pred :: State -> State
$ctoEnum :: Int -> State
toEnum :: Int -> State
$cfromEnum :: State -> Int
fromEnum :: State -> Int
$cenumFrom :: State -> [State]
enumFrom :: State -> [State]
$cenumFromThen :: State -> State -> [State]
enumFromThen :: State -> State -> [State]
$cenumFromTo :: State -> State -> [State]
enumFromTo :: State -> State -> [State]
$cenumFromThenTo :: State -> State -> State -> [State]
enumFromThenTo :: State -> State -> State -> [State]
Enum, State
State -> State -> Bounded State
forall a. a -> a -> Bounded a
$cminBound :: State
minBound :: State
$cmaxBound :: State
maxBound :: State
Bounded, Eq State
Eq State =>
(State -> State -> Ordering)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> State)
-> (State -> State -> State)
-> Ord State
State -> State -> Bool
State -> State -> Ordering
State -> State -> State
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: State -> State -> Ordering
compare :: State -> State -> Ordering
$c< :: State -> State -> Bool
< :: State -> State -> Bool
$c<= :: State -> State -> Bool
<= :: State -> State -> Bool
$c> :: State -> State -> Bool
> :: State -> State -> Bool
$c>= :: State -> State -> Bool
>= :: State -> State -> Bool
$cmax :: State -> State -> State
max :: State -> State -> State
$cmin :: State -> State -> State
min :: State -> State -> State
Ord)

-- | Move left or right; edges are clamped.
data Action = L | R
  deriving stock (Action -> Action -> Bool
(Action -> Action -> Bool)
-> (Action -> Action -> Bool) -> Eq Action
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Action -> Action -> Bool
== :: Action -> Action -> Bool
$c/= :: Action -> Action -> Bool
/= :: Action -> Action -> Bool
Eq, Int -> Action -> ShowS
[Action] -> ShowS
Action -> String
(Int -> Action -> ShowS)
-> (Action -> String) -> ([Action] -> ShowS) -> Show Action
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Action -> ShowS
showsPrec :: Int -> Action -> ShowS
$cshow :: Action -> String
show :: Action -> String
$cshowList :: [Action] -> ShowS
showList :: [Action] -> ShowS
Show, Int -> Action
Action -> Int
Action -> [Action]
Action -> Action
Action -> Action -> [Action]
Action -> Action -> Action -> [Action]
(Action -> Action)
-> (Action -> Action)
-> (Int -> Action)
-> (Action -> Int)
-> (Action -> [Action])
-> (Action -> Action -> [Action])
-> (Action -> Action -> [Action])
-> (Action -> Action -> Action -> [Action])
-> Enum Action
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Action -> Action
succ :: Action -> Action
$cpred :: Action -> Action
pred :: Action -> Action
$ctoEnum :: Int -> Action
toEnum :: Int -> Action
$cfromEnum :: Action -> Int
fromEnum :: Action -> Int
$cenumFrom :: Action -> [Action]
enumFrom :: Action -> [Action]
$cenumFromThen :: Action -> Action -> [Action]
enumFromThen :: Action -> Action -> [Action]
$cenumFromTo :: Action -> Action -> [Action]
enumFromTo :: Action -> Action -> [Action]
$cenumFromThenTo :: Action -> Action -> Action -> [Action]
enumFromThenTo :: Action -> Action -> Action -> [Action]
Enum, Action
Action -> Action -> Bounded Action
forall a. a -> a -> Bounded a
$cminBound :: Action
minBound :: Action
$cmaxBound :: Action
maxBound :: Action
Bounded, Eq Action
Eq Action =>
(Action -> Action -> Ordering)
-> (Action -> Action -> Bool)
-> (Action -> Action -> Bool)
-> (Action -> Action -> Bool)
-> (Action -> Action -> Bool)
-> (Action -> Action -> Action)
-> (Action -> Action -> Action)
-> Ord Action
Action -> Action -> Bool
Action -> Action -> Ordering
Action -> Action -> Action
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Action -> Action -> Ordering
compare :: Action -> Action -> Ordering
$c< :: Action -> Action -> Bool
< :: Action -> Action -> Bool
$c<= :: Action -> Action -> Bool
<= :: Action -> Action -> Bool
$c> :: Action -> Action -> Bool
> :: Action -> Action -> Bool
$c>= :: Action -> Action -> Bool
>= :: Action -> Action -> Bool
$cmax :: Action -> Action -> Action
max :: Action -> Action -> Action
$cmin :: Action -> Action -> Action
min :: Action -> Action -> Action
Ord)

-- | Deterministic transition.
--
-- >>> step R S0
-- S1
--
-- >>> step L S0
-- S0
--
-- >>> step R S2
-- Goal
step :: Action -> State -> State
step :: Action -> State -> State
step Action
L State
S0 = State
S0
step Action
L State
S1 = State
S0
step Action
L State
S2 = State
S1
step Action
L State
Goal = State
Goal
step Action
R State
S0 = State
S1
step Action
R State
S1 = State
S2
step Action
R State
S2 = State
Goal
step Action
R State
Goal = State
Goal

-- | State reward: living penalty, goal bonus.
--
-- >>> reward S0
-- -1.0
--
-- >>> reward Goal
-- 10.0
reward :: State -> Double
reward :: State -> Double
reward State
Goal = Double
10
reward State
_ = -Double
1

-- | One-step Bellman backup for a fixed deterministic policy.
bellmanPolicy :: Double -> Action -> (State -> Double) -> State -> Double
bellmanPolicy :: Double -> Action -> (State -> Double) -> State -> Double
bellmanPolicy Double
gamma Action
a State -> Double
v State
s = State -> Double
reward State
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
* State -> Double
v (Action -> State -> State
step Action
a State
s)

-- | One-step Bellman optimality backup.
bellmanOpt :: Double -> (State -> Double) -> State -> Double
bellmanOpt :: Double -> (State -> Double) -> State -> Double
bellmanOpt Double
gamma State -> Double
v State
s = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Double -> Action -> (State -> Double) -> State -> Double
bellmanPolicy Double
gamma Action
L State -> Double
v State
s) (Double -> Action -> (State -> Double) -> State -> Double
bellmanPolicy Double
gamma Action
R State -> Double
v State
s)

-- | Finite-horizon value iteration from the zero value function.
--
-- >>> valueIter 0 0.9 S0
-- 0.0
--
-- >>> valueIter 1 0.9 S0
-- -1.0
--
-- >>> valueIter 2 0.9 S0
-- -1.9
valueIter :: Int -> Double -> State -> Double
valueIter :: Int -> Double -> State -> Double
valueIter Int
0 Double
_ State
_ = Double
0
valueIter Int
n Double
gamma State
s = Double -> (State -> Double) -> State -> Double
bellmanOpt Double
gamma (Int -> Double -> State -> Double
valueIter (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double
gamma) State
s

-- | Greedy policy with respect to a value function.
optimalPolicy :: Double -> (State -> Double) -> State -> Action
optimalPolicy :: Double -> (State -> Double) -> State -> Action
optimalPolicy Double
gamma State -> Double
v State
s = (Action -> Action -> Ordering) -> [Action] -> Action
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy ((Action -> Double) -> Action -> Action -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (\Action
a -> Double -> Action -> (State -> Double) -> State -> Double
bellmanPolicy Double
gamma Action
a State -> Double
v State
s)) [Action
L, Action
R]

-- ---------------------------------------------------------------------------
-- Prob-composition view
-- ---------------------------------------------------------------------------

-- | State-dependent score modality. Not exported by 'Circuit.Prob' because it
-- leaks the input into the scalar map; useful for RL rewards.
scoreBy :: (a -> r -> r) -> Prob (->) r a a
scoreBy :: forall a r. (a -> r -> r) -> Prob (->) r a a
scoreBy a -> r -> r
f = (forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a)
-> (forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a
forall a b. (a -> b) -> a -> b
$ \(x, a) -> r
k (x
x, a
a) -> a -> r -> r
f a
a ((x, a) -> r
k (x
x, a
a))

-- | Transition as a Prob morphism.
transP :: Action -> Prob (->) Double State State
transP :: Action -> Prob (->) Double State State
transP = (State -> State) -> Prob (->) Double State State
forall a b r. (a -> b) -> Prob (->) r a b
embed ((State -> State) -> Prob (->) Double State State)
-> (Action -> State -> State)
-> Action
-> Prob (->) Double State State
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Action -> State -> State
step

-- | Reward as a state-dependent score modality.
rewardP :: Prob (->) Double State State
rewardP :: Prob (->) Double State State
rewardP = (State -> Double -> Double) -> Prob (->) Double State State
forall a r. (a -> r -> r) -> Prob (->) r a a
scoreBy (\State
s Double
v -> State -> Double
reward State
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
v)

-- | Discount as a scalar score modality.
discountP :: Double -> Prob (->) Double State State
discountP :: Double -> Prob (->) Double State State
discountP Double
gamma = (Double -> Double) -> Prob (->) Double State State
forall r a. (r -> r) -> Prob (->) r a a
score (Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
*)

-- | Bellman backup for a fixed action, expressed as three Prob morphisms:
-- reward, then discount, then transition. The contravariant composition in
-- 'Prob' reads right-to-left on continuations, so the written order is the
-- operational order.
bellmanP :: Double -> Action -> Prob (->) Double State State
bellmanP :: Double -> Action -> Prob (->) Double State State
bellmanP Double
gamma Action
a = Action -> Prob (->) Double State State
transP Action
a Prob (->) Double State State
-> Prob (->) Double State State -> Prob (->) Double State State
forall b c a.
Prob (->) Double b c
-> Prob (->) Double a b -> Prob (->) Double a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Double -> Prob (->) Double State State
discountP Double
gamma Prob (->) Double State State
-> Prob (->) Double State State -> Prob (->) Double State State
forall b c a.
Prob (->) Double b c
-> Prob (->) Double a b -> Prob (->) Double a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Prob (->) Double State State
rewardP

-- | Apply a Prob Bellman backup to a value function at a state.
backupP :: Double -> Action -> (State -> Double) -> State -> Double
backupP :: Double -> Action -> (State -> Double) -> State -> Double
backupP Double
gamma Action
a State -> Double
v State
s =
  Prob (->) Double State State
-> forall x. ((x, State) -> Double) -> (x, State) -> Double
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb (Double -> Action -> Prob (->) Double State State
bellmanP Double
gamma Action
a) (\(()
_, State
s') -> State -> Double
v State
s') ((), State
s)

-- ---------------------------------------------------------------------------
-- Discounted-return oracle
-- ---------------------------------------------------------------------------

-- | N-step discounted return via Prob composition.
--
-- Composes 'bellmanP gamma a' @n@ times via the 'Category' instance, then
-- applies the result to a zero continuation.  By the laws of 'Prob' 'Category'
-- composition, this is the n-step Bellman backup: reward on each step,
-- discounted and summed.
--
-- >>> discountedReturn 0.5 R 4 S0
-- -0.5
discountedReturn :: Double -> Action -> Int -> State -> Double
discountedReturn :: Double -> Action -> Int -> State -> Double
discountedReturn Double
gamma Action
a Int
n State
s =
  let chainP :: Prob (->) Double State State
chainP = (Prob (->) Double State State
 -> Prob (->) Double State State -> Prob (->) Double State State)
-> Prob (->) Double State State
-> [Prob (->) Double State State]
-> Prob (->) Double State State
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Prob (->) Double State State
-> Prob (->) Double State State -> Prob (->) Double State State
forall b c a.
Prob (->) Double b c
-> Prob (->) Double a b -> Prob (->) Double a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
(.) Prob (->) Double State State
forall a. Prob (->) Double a a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id (Int
-> Prob (->) Double State State -> [Prob (->) Double State State]
forall a. Int -> a -> [a]
replicate Int
n (Double -> Action -> Prob (->) Double State State
bellmanP Double
gamma Action
a))
   in Prob (->) Double State State
-> forall x. ((x, State) -> Double) -> (x, State) -> Double
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb Prob (->) Double State State
chainP (\((), State
_) -> Double
0) ((), State
s)

-- | Closed-form discounted return on a deterministic chain.
--
-- /Σ_{t=0}^{n-1}/ /γ/^t · reward(step^t(a, s)).
-- All terms are exact in 'Double' when /γ/ is a dyadic rational (e.g. 0.5)
-- and rewards are integers.
--
-- >>> closedFormReturn 0.5 R 4 S0
-- -0.5
closedFormReturn :: Double -> Action -> Int -> State -> Double
closedFormReturn :: Double -> Action -> Int -> State -> Double
closedFormReturn Double
gamma Action
a Int
n State
s =
  [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double) -> [Double] -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) ((State -> Double) -> [State] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map State -> Double
reward [State]
states) ((Double -> Double) -> Double -> [Double]
forall a. (a -> a) -> a -> [a]
iterate (Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
*) Double
1)
  where
    states :: [State]
states = Int -> [State] -> [State]
forall a. Int -> [a] -> [a]
take Int
n ([State] -> [State]) -> [State] -> [State]
forall a b. (a -> b) -> a -> b
$ (State -> State) -> State -> [State]
forall a. (a -> a) -> a -> [a]
iterate (Action -> State -> State
step Action
a) State
s

-- ---------------------------------------------------------------------------
-- Tropical / shortest-path row
-- ---------------------------------------------------------------------------

-- | Min-plus tropical semiring over 'Double'.
newtype Tropical = Tropical {Tropical -> Double
getTropical :: Double}
  deriving stock (Tropical -> Tropical -> Bool
(Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool) -> Eq Tropical
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Tropical -> Tropical -> Bool
== :: Tropical -> Tropical -> Bool
$c/= :: Tropical -> Tropical -> Bool
/= :: Tropical -> Tropical -> Bool
Eq, Eq Tropical
Eq Tropical =>
(Tropical -> Tropical -> Ordering)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Tropical)
-> (Tropical -> Tropical -> Tropical)
-> Ord Tropical
Tropical -> Tropical -> Bool
Tropical -> Tropical -> Ordering
Tropical -> Tropical -> Tropical
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Tropical -> Tropical -> Ordering
compare :: Tropical -> Tropical -> Ordering
$c< :: Tropical -> Tropical -> Bool
< :: Tropical -> Tropical -> Bool
$c<= :: Tropical -> Tropical -> Bool
<= :: Tropical -> Tropical -> Bool
$c> :: Tropical -> Tropical -> Bool
> :: Tropical -> Tropical -> Bool
$c>= :: Tropical -> Tropical -> Bool
>= :: Tropical -> Tropical -> Bool
$cmax :: Tropical -> Tropical -> Tropical
max :: Tropical -> Tropical -> Tropical
$cmin :: Tropical -> Tropical -> Tropical
min :: Tropical -> Tropical -> Tropical
Ord, Int -> Tropical -> ShowS
[Tropical] -> ShowS
Tropical -> String
(Int -> Tropical -> ShowS)
-> (Tropical -> String) -> ([Tropical] -> ShowS) -> Show Tropical
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Tropical -> ShowS
showsPrec :: Int -> Tropical -> ShowS
$cshow :: Tropical -> String
show :: Tropical -> String
$cshowList :: [Tropical] -> ShowS
showList :: [Tropical] -> ShowS
Show)

-- | Tropical addition: minimum.
tAdd :: Tropical -> Tropical -> Tropical
tAdd :: Tropical -> Tropical -> Tropical
tAdd (Tropical Double
a) (Tropical Double
b) = Double -> Tropical
Tropical (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
a Double
b)

-- | Tropical multiplication: ordinary addition.
tMul :: Tropical -> Tropical -> Tropical
tMul :: Tropical -> Tropical -> Tropical
tMul (Tropical Double
a) (Tropical Double
b) = Double -> Tropical
Tropical (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)

-- | Cost of each action: 0 at the goal, 1 elsewhere.
cost :: State -> Tropical
cost :: State -> Tropical
cost State
Goal = Double -> Tropical
Tropical Double
0
cost State
_ = Double -> Tropical
Tropical Double
1

-- | One-step tropical Bellman backup (shortest path to goal).
bellmanTropical :: (State -> Tropical) -> State -> Tropical
bellmanTropical :: (State -> Tropical) -> State -> Tropical
bellmanTropical State -> Tropical
v State
s =
  (Tropical -> Tropical -> Tropical) -> [Tropical] -> Tropical
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Tropical -> Tropical -> Tropical
tAdd [State -> Tropical
cost State
s Tropical -> Tropical -> Tropical
`tMul` State -> Tropical
v (Action -> State -> State
step Action
a State
s) | Action
a <- [Action
L, Action
R]]

-- | Finite-horizon shortest-path cost to goal.
--
-- >>> getTropical (shortestPath 0 S0)
-- Infinity
--
-- >>> getTropical (shortestPath 1 S0)
-- Infinity
--
-- >>> getTropical (shortestPath 2 S0)
-- Infinity
--
-- >>> getTropical (shortestPath 3 S0)
-- 3.0
--
-- >>> getTropical (shortestPath 3 Goal)
-- 0.0
shortestPath :: Int -> State -> Tropical
shortestPath :: Int -> State -> Tropical
shortestPath Int
0 State
Goal = Double -> Tropical
Tropical Double
0
shortestPath Int
0 State
_ = Double -> Tropical
Tropical (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
shortestPath Int
n State
s = (State -> Tropical) -> State -> Tropical
bellmanTropical (Int -> State -> Tropical
shortestPath (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) State
s

-- ---------------------------------------------------------------------------
-- System (Prob) view: controlled MDP
-- ---------------------------------------------------------------------------

-- | Local semiring class for the expectation runner (mirrors the kepler
-- executable, kept local to avoid substrate churn).
class Semiring r where
  sAdd :: r -> r -> r
  sMul :: r -> r -> r
  sZero :: r
  sOne :: r

instance Semiring Double where
  sAdd :: Double -> Double -> Double
sAdd = Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)
  sMul :: Double -> Double -> Double
sMul = Double -> Double -> Double
forall a. Num a => a -> a -> a
(*)
  sZero :: Double
sZero = Double
0
  sOne :: Double
sOne = Double
1

-- | Step a finite-state stochastic Moore machine by expectation, exactly as in
-- the circuits keystone, but specialised to 'Mono i o' with full state
-- observation.
expectSystem ::
  (Eq s, Semiring r) =>
  [s] ->
  System (Prob (->) r) s (Mono i o) ->
  [i] ->
  (s -> r) ->
  s ->
  r
expectSystem :: forall s r i o.
(Eq s, Semiring r) =>
[s]
-> System (Prob (->) r) s (Mono i o) -> [i] -> (s -> r) -> s -> r
expectSystem [s]
states System (Prob (->) r) s (Mono i o)
sys [i]
is s -> r
q s
s0 =
  (r -> r -> r) -> r -> [r] -> r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' r -> r -> r
forall r. Semiring r => r -> r -> r
sAdd r
forall r. Semiring r => r
sZero [s -> r
q s
s r -> r -> r
forall r. Semiring r => r -> r -> r
`sMul` s -> r
distFinal s
s | s
s <- [s]
states]
  where
    distFinal :: s -> r
distFinal = ((s -> r) -> i -> s -> r) -> (s -> r) -> [i] -> s -> r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (s -> r) -> i -> s -> r
stepDist s -> r
initDist [i]
is
    initDist :: s -> r
initDist s
s = if s
s s -> s -> Bool
forall a. Eq a => a -> a -> Bool
== s
s0 then r
forall r. Semiring r => r
sOne else r
forall r. Semiring r => r
sZero
    stepDist :: (s -> r) -> i -> s -> r
stepDist s -> r
dist i
i s
s' =
      (r -> r -> r) -> r -> [r] -> r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' r -> r -> r
forall r. Semiring r => r -> r -> r
sAdd r
forall r. Semiring r => r
sZero [s -> r
dist s
s r -> r -> r
forall r. Semiring r => r -> r -> r
`sMul` s -> i -> s -> r
pTrans s
s i
i s
s' | s
s <- [s]
states]
    pTrans :: s -> i -> s -> r
pTrans s
s i
i s
s' =
      Prob (->) r (s, Either Void i) (s, (o, ()))
-> forall x.
   ((x, (s, (o, ()))) -> r) -> (x, (s, Either Void i)) -> r
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb
        (System (Prob (->) r) s (Mono i o)
-> Prob (->) r (s, Dir (Mono i o)) (s, Pos (Mono i o))
forall (arr :: * -> * -> *) s (p :: Poly).
System arr s p -> arr (s, Dir p) (s, Pos p)
runSystem System (Prob (->) r) s (Mono i o)
sys)
        (\((), (s
s'', (o, ())
_)) -> if s
s' s -> s -> Bool
forall a. Eq a => a -> a -> Bool
== s
s'' then r
forall r. Semiring r => r
sOne else r
forall r. Semiring r => r
sZero)
        ((), (s
s, i -> Dir (Mono i (ZonkAny 0))
forall i o. i -> Dir (Mono i o)
monoIn i
i))

-- | The gridworld as a controlled stochastic Moore machine.
--
-- Input: action ('L' or 'R'). Output: full state observation.
gridSystem :: System (Prob (->) Double) State (Mono Action State)
gridSystem :: System (Prob (->) Double) State (Mono Action State)
gridSystem = Prob
  (->)
  Double
  (State, Dir (Mono Action State))
  (State, Pos (Mono Action State))
-> System (Prob (->) Double) State (Mono Action State)
forall (arr :: * -> * -> *) s (p :: Poly).
arr (s, Dir p) (s, Pos p) -> System arr s p
system (Prob
   (->)
   Double
   (State, Dir (Mono Action State))
   (State, Pos (Mono Action State))
 -> System (Prob (->) Double) State (Mono Action State))
-> Prob
     (->)
     Double
     (State, Dir (Mono Action State))
     (State, Pos (Mono Action State))
-> System (Prob (->) Double) State (Mono Action State)
forall a b. (a -> b) -> a -> b
$ (forall x.
 ((x, (State, Pos (Mono Action State))) -> Double)
 -> (x, (State, Dir (Mono Action State))) -> Double)
-> Prob
     (->)
     Double
     (State, Dir (Mono Action State))
     (State, Pos (Mono Action State))
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x.
  ((x, (State, Pos (Mono Action State))) -> Double)
  -> (x, (State, Dir (Mono Action State))) -> Double)
 -> Prob
      (->)
      Double
      (State, Dir (Mono Action State))
      (State, Pos (Mono Action State)))
-> (forall x.
    ((x, (State, Pos (Mono Action State))) -> Double)
    -> (x, (State, Dir (Mono Action State))) -> Double)
-> Prob
     (->)
     Double
     (State, Dir (Mono Action State))
     (State, Pos (Mono Action State))
forall a b. (a -> b) -> a -> b
$ \(x, (State, Pos (Mono Action State))) -> Double
k (x
x, (State
s, Dir (Mono Action State)
d)) ->
  let s' :: State
s' = Action -> State -> State
step (Dir (Mono Action (ZonkAny 1)) -> Action
forall i o. Dir (Mono i o) -> i
monoDir Dir (Mono Action (ZonkAny 1))
Dir (Mono Action State)
d) State
s
   in (x, (State, Pos (Mono Action State))) -> Double
k (x
x, (State
s', (State
s', ())))

-- ---------------------------------------------------------------------------
-- MDP and POMDP polynomial shapes
-- ---------------------------------------------------------------------------

-- | MDP interface: action in, next-state and reward out.
--
-- This matches the instance-table claim that the MDP row uses
-- @Mono a (s', r)@.  The reward is pinned on the current state to match
-- 'bellmanSystem' / 'bellmanOpt'.
mdpSystem :: System (Prob (->) Double) State (Mono Action (State, Double))
mdpSystem :: System (Prob (->) Double) State (Mono Action (State, Double))
mdpSystem = Prob
  (->)
  Double
  (State, Dir (Mono Action (State, Double)))
  (State, Pos (Mono Action (State, Double)))
-> System (Prob (->) Double) State (Mono Action (State, Double))
forall (arr :: * -> * -> *) s (p :: Poly).
arr (s, Dir p) (s, Pos p) -> System arr s p
system (Prob
   (->)
   Double
   (State, Dir (Mono Action (State, Double)))
   (State, Pos (Mono Action (State, Double)))
 -> System (Prob (->) Double) State (Mono Action (State, Double)))
-> Prob
     (->)
     Double
     (State, Dir (Mono Action (State, Double)))
     (State, Pos (Mono Action (State, Double)))
-> System (Prob (->) Double) State (Mono Action (State, Double))
forall a b. (a -> b) -> a -> b
$ (forall x.
 ((x, (State, Pos (Mono Action (State, Double)))) -> Double)
 -> (x, (State, Dir (Mono Action (State, Double)))) -> Double)
-> Prob
     (->)
     Double
     (State, Dir (Mono Action (State, Double)))
     (State, Pos (Mono Action (State, Double)))
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x.
  ((x, (State, Pos (Mono Action (State, Double)))) -> Double)
  -> (x, (State, Dir (Mono Action (State, Double)))) -> Double)
 -> Prob
      (->)
      Double
      (State, Dir (Mono Action (State, Double)))
      (State, Pos (Mono Action (State, Double))))
-> (forall x.
    ((x, (State, Pos (Mono Action (State, Double)))) -> Double)
    -> (x, (State, Dir (Mono Action (State, Double)))) -> Double)
-> Prob
     (->)
     Double
     (State, Dir (Mono Action (State, Double)))
     (State, Pos (Mono Action (State, Double)))
forall a b. (a -> b) -> a -> b
$ \(x, (State, Pos (Mono Action (State, Double)))) -> Double
k (x
x, (State
s, Dir (Mono Action (State, Double))
d)) ->
  let a :: Action
a = Dir (Mono Action (ZonkAny 2)) -> Action
forall i o. Dir (Mono i o) -> i
monoDir Dir (Mono Action (ZonkAny 2))
Dir (Mono Action (State, Double))
d
      s' :: State
s' = Action -> State -> State
step Action
a State
s
   in (x, (State, Pos (Mono Action (State, Double)))) -> Double
k (x
x, (State
s', ((State
s', State -> Double
reward State
s), ())))

-- | Check one deterministic MDP step by continuation.
mdpCheck :: Action -> State -> State -> Double -> Bool
mdpCheck :: Action -> State -> State -> Double -> Bool
mdpCheck Action
a State
s State
expectedS' Double
expectedR =
  Prob
  (->)
  Double
  (State, Either Void Action)
  (State, ((State, Double), ()))
-> forall x.
   ((x, (State, ((State, Double), ()))) -> Double)
   -> (x, (State, Either Void Action)) -> Double
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb (System (Prob (->) Double) State (Mono Action (State, Double))
-> Prob
     (->)
     Double
     (State, Dir (Mono Action (State, Double)))
     (State, Pos (Mono Action (State, Double)))
forall (arr :: * -> * -> *) s (p :: Poly).
System arr s p -> arr (s, Dir p) (s, Pos p)
runSystem System (Prob (->) Double) State (Mono Action (State, Double))
mdpSystem) ((), (State, ((State, Double), ()))) -> Double
checkCont ((), (State
s, Action -> Dir (Mono Action (ZonkAny 3))
forall i o. i -> Dir (Mono i o)
monoIn Action
a)) Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1.0
  where
    checkCont :: ((), (State, ((State, Double), ()))) -> Double
checkCont (()
_, (State
_sNext, ((State
s'', Double
r), ()))) =
      if State
s'' State -> State -> Bool
forall a. Eq a => a -> a -> Bool
== State
expectedS' Bool -> Bool -> Bool
&& Double
r Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
expectedR then Double
1.0 else Double
0.0

-- | POMDP observation: coarse location, not the true state.
data Observation = Far | Near | AtGoal
  deriving stock (Observation -> Observation -> Bool
(Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool) -> Eq Observation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Observation -> Observation -> Bool
== :: Observation -> Observation -> Bool
$c/= :: Observation -> Observation -> Bool
/= :: Observation -> Observation -> Bool
Eq, Int -> Observation -> ShowS
[Observation] -> ShowS
Observation -> String
(Int -> Observation -> ShowS)
-> (Observation -> String)
-> ([Observation] -> ShowS)
-> Show Observation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Observation -> ShowS
showsPrec :: Int -> Observation -> ShowS
$cshow :: Observation -> String
show :: Observation -> String
$cshowList :: [Observation] -> ShowS
showList :: [Observation] -> ShowS
Show, Int -> Observation
Observation -> Int
Observation -> [Observation]
Observation -> Observation
Observation -> Observation -> [Observation]
Observation -> Observation -> Observation -> [Observation]
(Observation -> Observation)
-> (Observation -> Observation)
-> (Int -> Observation)
-> (Observation -> Int)
-> (Observation -> [Observation])
-> (Observation -> Observation -> [Observation])
-> (Observation -> Observation -> [Observation])
-> (Observation -> Observation -> Observation -> [Observation])
-> Enum Observation
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Observation -> Observation
succ :: Observation -> Observation
$cpred :: Observation -> Observation
pred :: Observation -> Observation
$ctoEnum :: Int -> Observation
toEnum :: Int -> Observation
$cfromEnum :: Observation -> Int
fromEnum :: Observation -> Int
$cenumFrom :: Observation -> [Observation]
enumFrom :: Observation -> [Observation]
$cenumFromThen :: Observation -> Observation -> [Observation]
enumFromThen :: Observation -> Observation -> [Observation]
$cenumFromTo :: Observation -> Observation -> [Observation]
enumFromTo :: Observation -> Observation -> [Observation]
$cenumFromThenTo :: Observation -> Observation -> Observation -> [Observation]
enumFromThenTo :: Observation -> Observation -> Observation -> [Observation]
Enum, Observation
Observation -> Observation -> Bounded Observation
forall a. a -> a -> Bounded a
$cminBound :: Observation
minBound :: Observation
$cmaxBound :: Observation
maxBound :: Observation
Bounded, Eq Observation
Eq Observation =>
(Observation -> Observation -> Ordering)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Observation)
-> (Observation -> Observation -> Observation)
-> Ord Observation
Observation -> Observation -> Bool
Observation -> Observation -> Ordering
Observation -> Observation -> Observation
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Observation -> Observation -> Ordering
compare :: Observation -> Observation -> Ordering
$c< :: Observation -> Observation -> Bool
< :: Observation -> Observation -> Bool
$c<= :: Observation -> Observation -> Bool
<= :: Observation -> Observation -> Bool
$c> :: Observation -> Observation -> Bool
> :: Observation -> Observation -> Bool
$c>= :: Observation -> Observation -> Bool
>= :: Observation -> Observation -> Bool
$cmax :: Observation -> Observation -> Observation
max :: Observation -> Observation -> Observation
$cmin :: Observation -> Observation -> Observation
min :: Observation -> Observation -> Observation
Ord)

-- | Coarse observation function.
observe :: State -> Observation
observe :: State -> Observation
observe State
S0 = Observation
Far
observe State
S1 = Observation
Near
observe State
S2 = Observation
Near
observe State
Goal = Observation
AtGoal

-- | POMDP interface: hidden state carried as a 'Const' position, external loop
-- is action in / observation out.
--
-- This matches the instance-table claim that the POMDP row uses a state-hiding
-- @Prod (Const s) (Mono a o)@.  The @Const s@ position exposes the hidden
-- state as output but supplies no direction, so the external agent cannot feed
-- it back as input.
pomdpSystem :: System (Prob (->) Double) State (Prod (Const State) (Mono Action Observation))
pomdpSystem :: System
  (Prob (->) Double)
  State
  ('Prod ('Const State) (Mono Action Observation))
pomdpSystem = Prob
  (->)
  Double
  (State, Dir ('Prod ('Const State) (Mono Action Observation)))
  (State, Pos ('Prod ('Const State) (Mono Action Observation)))
-> System
     (Prob (->) Double)
     State
     ('Prod ('Const State) (Mono Action Observation))
forall (arr :: * -> * -> *) s (p :: Poly).
arr (s, Dir p) (s, Pos p) -> System arr s p
system (Prob
   (->)
   Double
   (State, Dir ('Prod ('Const State) (Mono Action Observation)))
   (State, Pos ('Prod ('Const State) (Mono Action Observation)))
 -> System
      (Prob (->) Double)
      State
      ('Prod ('Const State) (Mono Action Observation)))
-> Prob
     (->)
     Double
     (State, Dir ('Prod ('Const State) (Mono Action Observation)))
     (State, Pos ('Prod ('Const State) (Mono Action Observation)))
-> System
     (Prob (->) Double)
     State
     ('Prod ('Const State) (Mono Action Observation))
forall a b. (a -> b) -> a -> b
$ (forall x.
 ((x, (State, Pos ('Prod ('Const State) (Mono Action Observation))))
  -> Double)
 -> (x,
     (State, Dir ('Prod ('Const State) (Mono Action Observation))))
    -> Double)
-> Prob
     (->)
     Double
     (State, Dir ('Prod ('Const State) (Mono Action Observation)))
     (State, Pos ('Prod ('Const State) (Mono Action Observation)))
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x.
  ((x, (State, Pos ('Prod ('Const State) (Mono Action Observation))))
   -> Double)
  -> (x,
      (State, Dir ('Prod ('Const State) (Mono Action Observation))))
     -> Double)
 -> Prob
      (->)
      Double
      (State, Dir ('Prod ('Const State) (Mono Action Observation)))
      (State, Pos ('Prod ('Const State) (Mono Action Observation))))
-> (forall x.
    ((x, (State, Pos ('Prod ('Const State) (Mono Action Observation))))
     -> Double)
    -> (x,
        (State, Dir ('Prod ('Const State) (Mono Action Observation))))
       -> Double)
-> Prob
     (->)
     Double
     (State, Dir ('Prod ('Const State) (Mono Action Observation)))
     (State, Pos ('Prod ('Const State) (Mono Action Observation)))
forall a b. (a -> b) -> a -> b
$ \(x, (State, Pos ('Prod ('Const State) (Mono Action Observation))))
-> Double
k (x
x, (State
s, Dir ('Prod ('Const State) (Mono Action Observation))
d)) ->
  case Dir ('Prod ('Const State) (Mono Action Observation))
d of
    Left Void
v -> Void -> Double
forall a. Void -> a
absurd Void
v
    Right Either Void Action
dMono -> case Either Void Action
dMono of
      Left Void
v -> Void -> Double
forall a. Void -> a
absurd Void
v
      Right Action
a ->
        let s' :: State
s' = Action -> State -> State
step Action
a State
s
            o :: Observation
o = State -> Observation
observe State
s'
         in (x, (State, Pos ('Prod ('Const State) (Mono Action Observation))))
-> Double
k (x
x, (State
s', (State
s', (Observation
o, ()))))

-- | Check one deterministic POMDP step by continuation.
pomdpCheck :: Action -> State -> State -> Observation -> Bool
pomdpCheck :: Action -> State -> State -> Observation -> Bool
pomdpCheck Action
a State
s State
expectedS' Observation
expectedO =
  Prob
  (->)
  Double
  (State, Either Void (Either Void Action))
  (State, (State, (Observation, ())))
-> forall x.
   ((x, (State, (State, (Observation, ())))) -> Double)
   -> (x, (State, Either Void (Either Void Action))) -> Double
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb (System
  (Prob (->) Double)
  State
  ('Prod ('Const State) (Mono Action Observation))
-> Prob
     (->)
     Double
     (State, Dir ('Prod ('Const State) (Mono Action Observation)))
     (State, Pos ('Prod ('Const State) (Mono Action Observation)))
forall (arr :: * -> * -> *) s (p :: Poly).
System arr s p -> arr (s, Dir p) (s, Pos p)
runSystem System
  (Prob (->) Double)
  State
  ('Prod ('Const State) (Mono Action Observation))
pomdpSystem) ((), (State, (State, (Observation, ())))) -> Double
checkCont ((), (State
s, Either Void Action -> Either Void (Either Void Action)
forall a b. b -> Either a b
Right (Action -> Dir (Mono Action (ZonkAny 4))
forall i o. i -> Dir (Mono i o)
monoIn Action
a))) Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1.0
  where
    checkCont :: ((), (State, (State, (Observation, ())))) -> Double
checkCont (()
_, (State
_sNext, (State
hidden, (Observation
obs, ())))) =
      if State
hidden State -> State -> Bool
forall a. Eq a => a -> a -> Bool
== State
expectedS' Bool -> Bool -> Bool
&& Observation
obs Observation -> Observation -> Bool
forall a. Eq a => a -> a -> Bool
== Observation
expectedO then Double
1.0 else Double
0.0

-- | One-step Bellman optimality backup via 'System (Prob)'.
--
-- Reward is pinned on the /current/ state (matching 'bellmanOpt'); the System
-- runner computes the expected discounted future value of the next state.
bellmanSystem :: Double -> (State -> Double) -> State -> Double
bellmanSystem :: Double -> (State -> Double) -> State -> Double
bellmanSystem Double
gamma State -> Double
v State
s =
  State -> Double
reward State
s
    Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
gamma
      Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum
        [ [State]
-> System (Prob (->) Double) State (Mono Action State)
-> [Action]
-> (State -> Double)
-> State
-> Double
forall s r i o.
(Eq s, Semiring r) =>
[s]
-> System (Prob (->) r) s (Mono i o) -> [i] -> (s -> r) -> s -> r
expectSystem
            [State
S0, State
S1, State
S2, State
Goal]
            System (Prob (->) Double) State (Mono Action State)
gridSystem
            [Action
a]
            State -> Double
v
            State
s
        | Action
a <- [Action
L, Action
R]
        ]

-- | Finite-horizon value iteration using the 'System' runner.
valueIterSystem :: Int -> Double -> State -> Double
valueIterSystem :: Int -> Double -> State -> Double
valueIterSystem Int
0 Double
_ State
_ = Double
0
valueIterSystem Int
n Double
gamma State
s = Double -> (State -> Double) -> State -> Double
bellmanSystem Double
gamma (Int -> Double -> State -> Double
valueIterSystem (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double
gamma) State
s