Marginalising the MDP Dynamics Function

Problem Statement

An MDP is fully specified by one four-argument function $p(s', r \mid s, a)$. Derive the two quantities every algorithm actually consumes: the expected immediate reward $r(s,a)$ and the state-transition probabilities $p(s' \mid s, a)$.

Background

The five-tuple definition of an MDP is $\langle \mathcal{S}, \mathcal{A}, p, r, \gamma \rangle$, and the single object that carries all the dynamics is

$p(s', r \mid s, a) = \Pr\{S_{t+1}=s', R_{t+1}=r \mid S_t=s, A_t=a\}$

a joint distribution over the pair (next state, reward). Almost no algorithm uses it in that form. Bellman backups need two marginals of it:

$r(s,a) = \sum_{s'}\sum_{r} r\, p(s', r \mid s, a) \qquad p(s' \mid s, a) = \sum_{r} p(s', r \mid s, a)$

The subtlety is that the same successor state can be reached with different rewards. When you marginalise out the reward you must accumulate those probabilities, not overwrite them. A stochastic reward on a fixed transition is common (a noisy sensor, a randomised bonus), and an implementation that keeps only the last entry it saw produces a transition matrix whose rows do not sum to 1 — after which every value estimate downstream is quietly wrong.

Your Task

Implement:

def expected_reward(dynamics, s, a):
    ...

def transition_probs(dynamics, s, a, n_states):
    ...

dynamics is a dict mapping the key (s, a) to a list of (s_next, reward, prob) tuples. The probabilities for one (s, a) key sum to 1.

• expected_reward returns a float, the expected immediate reward r(s, a).
• transition_probs returns a list of n_states floats, entry s2 being p(s2 | s, a).

Input / Output Format

Input is a dict keyed by (state, action) tuples. Outputs are a float and a list of floats, rounded to 4 decimals by the grader.

Sample

dyn = {(0, 0): [(0, 1.0, 0.3), (1, 1.0, 0.2), (1, -1.0, 0.4), (0, 2.0, 0.1)]}
print(round(expected_reward(dyn, 0, 0), 4))
print([round(p, 4) for p in transition_probs(dyn, 0, 0, 3)])

Output:

0.3
[0.4, 0.6, 0.0]

Expected reward: 0.31 + 0.21 + 0.4(-1) + 0.12 = 0.3**. State 0 is reached by two entries (0.3 + 0.1 = 0.4), state 1 by two entries (0.2 + 0.4 = 0.6).