🤖 machine learning

Distributions as Actions: A Unified Framework for Diverse Action Spaces

This paper introduces a unified reinforcement learning framework that redefines actions as parameterized distributions to create a continuous action space, enabling a low-variance gradient estimator and an effective actor-critic algorithm (DA-AC) that achieves competitive performance across discrete, continuous, and hybrid control tasks.

Original authors: Jiamin He, A. Rupam Mahmood, Martha White

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiamin He, A. Rupam Mahmood, Martha White

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to play a video game. In the world of Reinforcement Learning (RL), the robot is the "agent," and the game world is the "environment." Every time the robot makes a move, it gets a reward (points) or a penalty. The goal is to learn the best moves to get the most points.

Usually, the way we teach these robots depends entirely on what kind of moves they can make:

If the game has discrete moves (like pressing "Left," "Right," or "Jump"), we use one type of math.
If the game has continuous moves (like steering a car with infinite precision), we use a different type of math.
If the game has a mix of both, we have to build a complicated, custom solution for that specific game.

This paper introduces a new way of thinking called "Distributions as Actions." It's like finding a universal translator that lets us use the same simple math for all types of games, whether they are discrete, continuous, or a mix.

Here is how it works, using some everyday analogies:

1. The Big Shift: Changing the "Handshake"

In the old way (Classical RL), the robot decides exactly what to do.

Analogy: Imagine a chef (the robot) telling the waiter (the environment) exactly what to serve: "Bring me the red soup." The waiter just brings it.

In this new framework (Distributions as Actions), the robot stops deciding the exact move. Instead, it decides the recipe for the move.

Analogy: Now, the chef tells the waiter: "Bring me a soup that is 70% red and 30% blue." The waiter then randomly picks a bowl based on those instructions.
The Magic: Even if the final soup is either "Red" or "Blue" (discrete), the chef's instruction ("70/30") is a smooth, continuous number. This allows the researchers to treat every problem as if it were a smooth, continuous one, even if the actual moves are choppy or discrete.

2. The New Gradient: Smoother Roads

To learn, the robot needs to know which direction to nudge its "recipe" to get better scores. This is called a "gradient."

The Problem: In the old way, if the robot is guessing discrete moves, the math is very "noisy" (like trying to drive on a bumpy dirt road). The robot takes big, shaky steps and learns slowly.
The Solution (DA-PG): Because the robot is now adjusting smooth numbers (the recipe percentages) instead of guessing specific moves, the math becomes much smoother (like driving on a paved highway).
The Result: The robot learns faster and more stably because the "road" it travels on has fewer bumps. The paper proves mathematically that this new method has less "noise" (variance) than the old methods.

3. The Critic's Dilemma: The Judge Needs a Map

In these systems, there is a "Critic" (a judge) that tells the robot how good a move was.

The Challenge: Because the robot is now outputting "recipes" (distributions) instead of specific moves, the judge has to evaluate a whole range of possibilities at once. This is harder for the judge to learn. It's like a judge who used to grade a single essay now has to grade a whole library of potential essays.
The Fix (ICL): The authors invented a trick called Interpolated Critic Learning.
- Analogy: Imagine the judge is trying to learn the best recipe. Instead of only looking at the exact recipe the robot used, the judge also looks at "in-between" recipes. If the robot used a "70/30" mix, the judge also checks a "60/40" mix and an "80/20" mix.
- Why it helps: This forces the judge to learn the shape of the landscape, not just the specific points. It helps the judge understand that "moving slightly toward 80/20 might be better," giving the robot a clearer map on how to improve.

4. The Result: One Algorithm to Rule Them All

The authors built a robot algorithm called DA-AC (Distributions-as-Actions Actor-Critic) using these ideas. They tested it in three very different worlds:

Continuous: Like controlling a robot arm with infinite precision.
Discrete: Like playing a game with only "Up, Down, Left, Right."
Hybrid: Like a game where you choose a weapon (discrete) and then aim it (continuous).

The Claim: In all these different worlds, DA-AC performed just as well as, or better than, the specialized algorithms that were previously designed just for those specific worlds.

Summary

The paper argues that by redefining the boundary between the "agent" and the "environment"—treating the probability settings of a move as the move itself—we can unify reinforcement learning.

Old Way: Different tools for different jobs (a hammer for nails, a screwdriver for screws).
New Way: A universal multi-tool that works on nails, screws, and bolts equally well, because it changes the way it interacts with the object to make the job easier.

The authors show that this "universal multi-tool" (DA-AC) is not just a theoretical idea; it actually works better in practice across a wide variety of complex tasks.

Technical Summary: Distributions as Actions: A Unified Framework for Diverse Action Spaces

Problem Statement

Reinforcement learning (RL) algorithms are traditionally categorized by their suitability for specific action space structures. Value-based methods (e.g., DQN) excel in discrete spaces by enumerating action values, while policy-based methods (e.g., DPG, SAC) are designed for continuous spaces, leveraging the differentiability of action-value functions. However, current algorithms remain tightly coupled to the specific structure of the action space. Designing unified algorithms that generalize across discrete, continuous, and hybrid action spaces is difficult because different gradient estimators and architectures are often required for each domain. Furthermore, existing policy gradient estimators for stochastic policies, such as the Likelihood-Ratio (LR) and Reparameterization (RP) methods, often suffer from high variance or introduce bias, particularly in discrete settings.

Methodology

The Distributions-as-Actions (DA) Framework

The paper introduces a novel formulation that redefines the boundary between the agent and the environment. In the classical RL framework, the agent's policy $\pi_\theta$ outputs an action (or distribution parameters) and samples the action internally. In the proposed Distributions-as-Actions (DA) framework:

The agent outputs distribution parameters $\bar{\pi}_\theta(s)$ (e.g., mean and variance for Gaussian, probability vectors for categorical) as its action.
The sampling process $A \sim f(\cdot|\bar{\pi}_\theta(s))$ is treated as part of the environment's stochastic transition function.

This transformation converts the original action space $\mathcal{A}$ (which may be discrete, continuous, or hybrid) into a continuous distribution parameter space $\mathcal{U}$ . Consequently, the agent operates within a new Markov Decision Process (DA-MDP) with a continuous action space $\mathcal{U}$ , regardless of the underlying nature of $\mathcal{A}$ .

Distributions-as-Actions Policy Gradient (DA-PG)

Building on the DA-MDP, the authors derive the DA-PG estimator. By applying the Deterministic Policy Gradient (DPG) theorem to the continuous parameter space $\mathcal{U}$ , they obtain a gradient estimator:
$\hat{\nabla}_\theta \hat{J}(\bar{\pi}_\theta; S_t) = \nabla_\theta \bar{\pi}_\theta(S_t)^\top \nabla_U \bar{Q}_w(S_t, U)|_{U=\bar{\pi}_\theta(S_t)}$
where $\bar{Q}_w$ is a critic estimating the expected return over the action distribution defined by parameters $U$ .

Theoretical Properties:

Variance Reduction: The paper proves that DA-PG is the conditional expectation of both the LR and RP estimators. By the law of total variance, this results in strictly lower variance than the original estimators, provided the action-conditioned or noise-induced variance is non-zero.
Unification: When the distribution is a Dirac delta, DA-PG reduces exactly to the standard DPG estimator, establishing it as a strict generalization.

Interpolated Critic Learning (ICL)

A challenge in the DA framework is that the critic $\bar{Q}_w$ must learn over a higher-dimensional input space (distribution parameters), which can be difficult and lead to increased bias if the critic is inaccurate. To address this, the authors propose Interpolated Critic Learning (ICL).

Mechanism: Instead of updating the critic only at the sampled distribution parameters $U_t$ , ICL utilizes the observed action $A_t$ to infer that the transition could have been generated by other distribution parameters.
Implementation: The critic is updated at an interpolated point $\hat{U}_t = \omega_t U_t + (1-\omega_t) U_{A_t}$ , where $U_{A_t}$ represents the deterministic parameters corresponding to the observed action $A_t$ , and $\omega_t \sim \text{Uniform}[0, 1]$ .
Goal: This encourages the critic to learn a smoother value landscape with informative curvature, facilitating better policy optimization in the parameter space.

DA-AC Algorithm

The authors implement a practical deep RL algorithm, Distributions-as-Actions Actor-Critic (DA-AC), based on the TD3 architecture. DA-AC incorporates the DA-PG estimator for the actor and the ICL update for the critic, enabling a single algorithm to handle diverse action spaces.

Key Contributions

Unified Framework: A reparameterization of the RL problem that treats distribution parameters as actions, unifying discrete, continuous, and hybrid control under a single continuous-action algorithm.
DA-PG Estimator: A generalized policy gradient estimator with theoretically lower variance than standard LR and RP estimators, applicable to arbitrary action space types.
ICL Strategy: A simple yet effective technique to improve critic learning in the DA framework by leveraging off-distribution updates via interpolation.
Empirical Validation: Comprehensive evaluation showing that DA-AC achieves competitive or superior performance across continuous, discrete, and hybrid control benchmarks compared to specialized baselines.

Experimental Results

The paper evaluates DA-AC across four distinct settings:

Continuous Control: On MuJoCo and DeepMind Control Suite tasks, DA-AC achieves competitive performance, often outperforming TD3 and other baselines, particularly in high-dimensional environments.
Discrete Control: On Gym classic control and MinAtar environments, DA-AC performs comparably to DQN and other strong discrete baselines.
High-Dimensional Discrete Control: In discretized versions of continuous control tasks (e.g., Humanoid-v4 with 17 dimensions discretized into 7 bins), DA-AC significantly outperforms baselines like ST-AC, LR-AC, and PPO, where other methods struggle with the combinatorial explosion or high variance.
Hybrid Control: On Parameterized Action MDPs (PAMDPs), DA-AC achieves competitive results against specialized algorithms like PATD3 and PDQN.

Ablation Studies:
Experiments in bandit settings and high-dimensional discrete control demonstrate that removing ICL (DA-AC w/o ICL) leads to significantly worse performance, confirming that ICL is crucial for learning an effective critic in the DA framework.

Significance and Claims

The paper claims that the Distributions-as-Actions framework offers a principled way to decouple the learning algorithm from the specific structure of the action space. By shifting the agent-environment boundary, it allows for the application of low-variance, deterministic-style gradient methods to discrete and hybrid problems. The authors emphasize that this approach:

Provides the first unbiased, low-variance reparameterization-style estimator for discrete distributions.
Enables a single algorithmic framework (DA-AC) to handle a wide variety of control settings without needing domain-specific architectural changes.
Demonstrates that treating sampling as part of the environment can simplify learning and open new algorithmic perspectives, particularly for complex hybrid action spaces.

The work suggests that future research could extend this framework to model-based methods, hierarchical control, and more advanced critic learning strategies, but the current contribution focuses on establishing the viability of the framework through the DA-AC algorithm and its empirical success.