Marginalising the MDP Dynamics Function

easyFREEMarkov Decision Processes and Bellman EquationsReinforcement LearningRL 1: Foundations

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).

Example

Input

dyn = {(0, 0): [(0, 1.0, 0.3), (1, 1.0, 0.2), (1, -1.0, 0.4), (0, 2.0, 0.1)]}
expected_reward(dyn, 0, 0), transition_probs(dyn, 0, 0, 3)

Output

0.3
[0.4, 0.6, 0.0]

Reasoning

r(s,a) sums reward*prob over every entry: 0.3 + 0.2 - 0.4 + 0.2 = 0.3. p(s'|s,a) sums probability over rewards, so the two entries landing in state 0 combine to 0.4 and the two landing in state 1 combine to 0.6.

Constraints

1 <= n_states <= 100
The list for a given (s, a) may contain the same s_next more than once with different rewards.
Probabilities for one (s, a) sum to 1.
States not reachable from (s, a) must appear as 0.0 in the returned list.
Do not round inside the functions.

Learn this in

Reinforcement LearningML Study Plan Introduction to Reinforcement LearningReinforcement Learning Markov Decision ProcessesReinforcement Learning

More from RL 1: Foundations

GAE Over a Full Trajectoryhard PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium

Concept primer

1. Background Knowledge

In Markov Decision Processes (MDPs), the environment dynamics are fundamentally defined by the joint probability distribution $p(s', r \mid s, a)$ . This function describes the likelihood of transitioning to a specific next state $s'$ and receiving a specific reward $r$ , given that the agent is currently in state $s$ and takes action $a$ . While this joint distribution contains all necessary information, most reinforcement learning algorithms (such as Value Iteration or Policy Iteration) do not operate on the joint distribution directly. Instead, they require two marginalized quantities derived from it.

The first quantity is the expected immediate reward, denoted as $r(s, a)$ . This is a scalar value representing the average reward the agent expects to receive upon taking action $a$ in state $s$ . It is calculated by summing the products of each possible reward and its corresponding probability, marginalizing over both the next state and the reward variables. The second quantity is the state-transition probability, denoted as $p(s' \mid s, a)$ . This is a probability distribution over the next states, indicating the likelihood of ending up in any specific state $s'$ after taking action $a$ . It is obtained by summing the joint probabilities over all possible rewards for a given next state.

Understanding marginalization is critical here. The problem statement highlights a common implementation error: assuming a deterministic mapping from $(s, a)$ to $s'$ . In reality, a single transition $(s, a) \to s'$ can be associated with multiple different rewards, each with its own probability. When calculating $p(s' \mid s, a)$ , you must accumulate the probabilities of all entries that lead to the same $s'$ , regardless of the reward value. Failing to do so results in a transition matrix where rows do not sum to 1, violating the axioms of probability and causing downstream algorithmic failures.

2. Algorithm Approach

The core algorithmic pattern for this problem is aggregation via iteration. Since the input dynamics is provided as a dictionary of lists, you must iterate through the list of outcomes associated with the key (s, a).

For expected_reward, the approach is a straightforward weighted sum. You iterate through each tuple (s_next, reward, prob) in the list and accumulate reward * prob. This directly implements the definition of expected value for a discrete random variable.

For transition_probs, the approach is frequency accumulation. You need to initialize an array (or list) of size n_states with zeros. Then, iterate through the same list of outcomes. For each tuple (s_next, reward, prob), you add prob to the index s_next in your array. This effectively sums the probabilities for all paths leading to s_next, marginalizing out the reward variable. This is a classic bucketing or histogram accumulation pattern.

3. Step-by-Step Strategy

Extract the Outcome List: Access the list of tuples from the dynamics dictionary using the key (s, a). If the key does not exist, handle it appropriately (though the problem implies valid inputs).
Implement expected_reward:

Initialize a variable total_reward to 0.0.
Iterate through each (s_next, reward, prob) in the outcome list.
Update total_reward by adding reward * prob.
Return total_reward.

Implement transition_probs:

Initialize a list probs of length n_states with all zeros.
Iterate through each (s_next, reward, prob) in the outcome list.
Identify the index s_next.
Add prob to probs[s_next]. Note that you are adding, not assigning, because multiple entries may map to the same s_next.
Return the probs list.

Verification: Ensure that the sum of the returned probs list is approximately 1.0 (within floating-point tolerance). This serves as a sanity check for your marginalization logic.

4. Common Pitfalls

Overwriting vs. Accumulating: The most frequent error in transition_probs is using assignment (probs[s_next] = prob) instead of addition (probs[s_next] += prob). This fails when multiple outcomes lead to the same state with different rewards. Always use accumulation.
Index Out of Bounds: Ensure that s_next is a valid index for the probs list. The problem states n_states is provided, so initialize the list with exactly that size. If s_next >= n_states, it indicates an input inconsistency, but typically inputs are well-formed.
Floating Point Precision: While the grader rounds to 4 decimals, intermediate calculations should use standard floating-point arithmetic. Do not round intermediate values, as this can introduce cumulative errors.
Ignoring the Reward in Transition Probabilities: When calculating $p(s' \mid s, a)$ , the reward value in the tuple is irrelevant. Do not let the presence of the reward variable confuse the logic; simply sum the probabilities associated with the target state.
Empty Dynamics: Although unlikely in this specific problem context, consider what happens if the outcome list is empty. The expected reward should be 0, and transition probs should be all zeros.

5. Time & Space Complexity

Time Complexity:
expected_reward: $O(K)$ , where $K$ is the number of outcomes (tuples) for the given (s, a) pair. You iterate through the list once.
transition_probs: $O(K + N)$ , where $K$ is the number of outcomes and $N$ is n_states. Initializing the list takes $O(N)$ , and iterating through the outcomes takes $O(K)$ . Since $K$ is typically much smaller than $N$ in sparse MDPs, this is efficient.
Space Complexity:
expected_reward: $O(1)$ auxiliary space, as you only store a single accumulator variable.
transition_probs: $O(N)$ auxiliary space, required to store the output list of size n_states.

The overall efficiency is linear with respect to the number of possible outcomes for a given state-action pair, making it highly scalable for large MDPs where the dynamics are sparse.

solution.py

Python 3.13

Test results