Marginalize Next-State Distribution

mediumPROMarkov Decision Processes and Bellman EquationsReinforcement LearningRL 1: Foundations

Problem Statement

Given a policy pi over actions at a state, and a transition matrix P where P[a][s'] is the probability of reaching next state s' when taking action a, compute the marginal next-state distribution under the policy:

$P^\pi(s'\mid s) = \sum_a \pi(a\mid s)\, P(s'\mid s, a)$

Implement marginal_next_state(pi, P) returning a distribution over next states.

Example

Input

marginal_next_state([0.5, 0.5], [[1.0, 0.0], [0.0, 1.0]])

Output

[0.5, 0.5]

Reasoning

Initialize the next-state distribution vector to zeros, representing that no probability mass has been accumulated yet for any state.
Process the first action with policy probability $\pi(0) = 0.5$ : multiply this weight by the transition row $[1.0, 0.0]$ to get $[0.5, 0.0]$ , and add this to the accumulator, resulting in $[0.5, 0.0]$ .
Process the second action with policy probability $\pi(1) = 0.5$ : multiply this weight by the transition row $[0.0, 1.0]$ to get $[0.0, 0.5]$ , and add this to the current accumulator $[0.5, 0.0]$ .
Sum the contributions from both actions to obtain the final marginal probabilities: $0.5 + 0.0 = 0.5$ for the first state and $0.0 + 0.5 = 0.5$ for the second state.
The final output is [0.5, 0.5]

Constraints

len(pi) == len(P) (one row per action); all rows same length S.
Each P[a] sums to 1; pi sums to 1.
Output length S, sums to 1.

Learn this in

Probability & StatisticsFoundations 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 primerPREVIEW

1. Background Knowledge

In a Markov Decision Process (MDP), the environment is described by states, actions, and a transition function $P(s' \mid s, a)$ that gives the probability of transitioning to state $s'$ from state $s$ when taking action $a$ . A policy $\pi(a \mid s)$ is a probability distribution over actions at each state, representing the agent's strategy. When the agent follows a stochastic policy, the exact next state is not determined by a single action but by a mixture of all possible actions, weighted by their selection probabilities.

The marginal next-state distribution $P^\pi(s' \mid s)$ answers the question: "If I am in state $s$ and follow policy $\pi$ , what is the probability of ending up in each next state $s'$ ?" This is computed by summing over all actions, weighting each action's transition probability by the probability of choosing that action. This operation is fundamental in Bellman equations, where the value of a state under a policy depends on the expected value of next states, which in turn depends on this marginal distribution.

Mathematically, this is a weighted sum (or expectation) over the action space. If $\pi$ is a row vector of shape $(A,)$ and $P$ is a matrix of shape $(A, S')$ , the result is a row vector of shape $(S',)$ obtained by multiplying $\pi$ with $P$ . This is a standard matrix-vector multiplication in linear algebra, which is the computational backbone of many RL algorithms.

2. Algorithm Approach

The core approach is to compute the expectation of the next-state distribution over the action distribution. For each possible next state $s'$ , you sum the product of the policy probability for each action $a$ and the transition probability from $a$ to $s'$ .

In code, this translates to:

Iterate over each action $a$ .
For each action, scale the entire row $P[a]$ (which is a distribution over next states) by $\pi[a]$ .
Accumulate these scaled distributions into a single result vector.

This is equivalent to performing a dot product between the policy vector and the transition matrix. If you are using NumPy, this can be done in a single vectorized operation. If you are implementing it from scratch, you will use nested loops or a reduction operation.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results