Моделирование наблюдения для вероятностного языка программирования со свободными монадами

Я пытаюсь реализовать DST для вероятностного программирования, вдохновленный сообщением в блоге Джареда Тобина.

Бесплатная монада для DST, включая дистрибутив Бернулли и бета-версию, может выглядеть так:

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  deriving (Functor)

type Model = Free ModelF

bernoulli :: Double -> Model Bool
bernoulli p = liftF (BernoulliF p id)

beta :: Double -> Double -> Model Double
beta a b = liftF (BetaF a b id)

Я хотел бы расширить его, чтобы поддерживать наблюдение за дистрибутивами. То есть: принудительно указать определенное возвращаемое значение.

Первой попыткой было расширить функтор:

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  | ObsF (ModelF r) r
  deriving (Functor)

Однако это проблематично, поскольку допускает рекурсивное вложение ObsF.

Одним из решений было бы поднять свободную монаду в другую свободную монаду, представляющую наблюдаемое значение:

data ObsModelF r =
    PureF (Model r)
  | ObsF (Model r) r
  deriving (Functor)

observe :: Model r -> r -> Free ObsModelF r
observe model o = liftF (ObsF model o)

Синтаксис, который мне нужен, таков:

let dist = do
      p <- beta 1 1
      observe (bernoulli p) True

Это несовместимо с ObsModelF, так как beta 1 1 также должно быть преобразовано в ObsModelF. Для этого требуются два отдельных конструктора: один для распределений без наблюдения и один с наблюдением. Можно ли избежать этой сложности в DST?

Третья попытка. На этот раз я заключаю функтор распределения (ModelF) в функтор наблюдения, который затем поднимаю (ObsModelF). Я получаю некоторые ошибки типа, которые я не совсем понимаю (см. комментарии):

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  deriving (Functor)

data ObsModelF r =
    PureF (ModelF r)
  | ObsF (ModelF r) r
  deriving (Functor)

type ObsModel = Free ObsModelF

bernoulli :: Double -> ObsModel Bool
bernoulli p = liftF . PureF $ BernoulliF p id

beta :: Double -> Double -> ObsModel Double
beta a b = liftF . PureF $ BetaF a b id

observe :: ObsModel r -> r -> ObsModel r
observe f o = liftF $ case f of
           (Free f') -> case f' of
               (PureF f'') ->  ObsF f'' o   -----| Couldn't match expected type ‘Free ObsModelF r’
               (ObsF f'' o') ->  ObsF f'' o    --| with actual type ‘r’

toSampler :: (RandomGen g) => ObsModel r -> State g r
toSampler = iterM $ \case
  ObsF a o -> case a of 
    BernoulliF p f -> return o >>= f  -----| Couldn't match type ‘StateT g Data.Functor.Identity.Identity r’
    BetaF a b f    -> return o >>= f     --| with ‘Bool’
  PureF a -> case a of 
    BernoulliF p f -> MWC.bernoulli p >>= f
    BetaF a b f    -> MWC.beta a b >>= f

Я включил очень простой интерпретатор (toSampler). В настоящее время он полностью игнорирует распределение и просто плоско отображает значение наблюдения. Обратите внимание, что это не является предполагаемым поведением при обработке дистрибутива. Он предназначен только для того, чтобы показать, как интерпретатор будет проходить свободную структуру.