Softmax Policy Probabilities

easyPROReturns, Policies, and Value FunctionsReinforcement LearningRL 1: Foundations

Problem Statement

Convert action preferences h into a policy using the softmax (Boltzmann) distribution with temperature tau:

$\pi(a) = \frac{e^{h_a / \tau}}{\sum_b e^{h_b / \tau}}$

Implement softmax_policy(h, tau) returning a list of probabilities. Use the max-subtraction trick for numerical stability.

Example

Input

softmax_policy([1.0, 1.0], 1.0)

Output

[0.5, 0.5]

Reasoning

Divide each preference by the temperature to obtain the scaled logits: $z = [1.0/1.0, 1.0/1.0] = [1.0, 1.0]$ .
Identify the maximum value in the scaled logits for numerical stability: $m = \max(1.0, 1.0) = 1.0$ .
Compute the exponentials of the shifted values ( $z_i - m$ ) to prevent overflow: $e^{1.0 - 1.0} = e^0 = 1.0$ for both elements, resulting in $[1.0, 1.0]$ .
Sum the exponentials to determine the normalization constant: $s = 1.0 + 1.0 = 2.0$ .
Normalize each exponential by the sum to get the final probabilities: $1.0 / 2.0 = 0.5$ for each action.
The final output is [0.5, 0.5]

Constraints

1 <= len(h) <= 1000, tau > 0
Probabilities must sum to 1.
Subtract max(h/tau) before exponentiating for stability.

Learn this in

Policy Gradient MethodsReinforcement 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 reinforcement learning, an agent must choose actions based on its current state. A policy $\pi(a|s)$ defines the probability distribution over actions. When we have a vector of raw scores or preferences $h = [h_1, h_2, \dots, h_n]$ (often called logits or Q-values), we need a mechanism to convert these unbounded real numbers into valid probabilities that sum to 1. The softmax function (also known as the Boltzmann distribution) is the standard tool for this conversion. It maps any real-valued vector to a probability simplex, ensuring all outputs are positive and sum to unity.

The parameter $\tau$ (temperature) controls the "sharpness" of the distribution. When $\tau$ is small, the softmax becomes more peaked, favoring the action with the highest preference almost exclusively. When $\tau$ is large, the distribution becomes flatter, approaching a uniform distribution. This temperature parameter is crucial in exploration-exploitation trade-offs: low temperature encourages exploitation of known best actions, while high temperature encourages exploration.

A critical practical concern is numerical stability. Directly computing exponentials of large numbers can lead to overflow in floating-point arithmetic. For example, $e^{1000}$ exceeds the range of standard 64-bit floats. The max-subtraction trick addresses this by exploiting the mathematical property that softmax is invariant to adding a constant to all inputs: $\text{softmax}(h) = \text{softmax}(h - c)$ for any constant $c$ . By subtracting the maximum value before exponentiation, we ensure all exponents are non-positive, preventing overflow while preserving the final probability distribution.

2. Algorithm Approach

The solution follows a straightforward three-phase pattern:

Stabilize: Find the maximum value in the preference vector and subtract it from all elements. This ensures the largest exponent will be $e^0 = 1$ , keeping all exponential values in the range $(0, 1]$ .
Exponentiate: Apply the exponential function to each stabilized, temperature-scaled value.
Normalize: Sum all exponential values and divide each by this sum to produce valid probabilities.

This approach is O(n) in both time and space, making it efficient for typical action spaces. The key insight is that the max-subtraction step is mathematically equivalent to the original softmax but numerically robust.

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