RN-D: Discretized Categorical Actors for On-Policy Reinforcement Learning

AI summary

Overview

Research area: On-policy reinforcement learning for continuous control, specifically actor (policy) parameterization and actor-network architecture.

Technical level: Advanced. The paper combines a gradient-variance proposition, a loss-function reinterpretation of PPO, and an architectural design study across multiple benchmarks. The prose is readable, but the content assumes familiarity with PPO, policy gradients, and Gaussian policies.

Scope: The paper proposes RN-D, a regularized residual-network actor that outputs discretized categorical (bin-wise softmax) action distributions instead of the standard diagonal-Gaussian continuous actor, and evaluates it under PPO, TRPO, and SPO on Gym locomotion and ManiSkill benchmarks.

What This Paper Is About

Most on-policy RL systems for continuous control default to a Gaussian actor with a shallow MLP backbone, a choice made largely for analytical and computational convenience rather than because it optimizes well. The authors argue that this parameterization is an underexplored design lever, and that it makes the actor update behave like a weighted squared-error regression problem, which is fragile when gradients are noisy and updates must be conservative. Their goal is to test whether replacing it with a factorized categorical policy over uniformly discretized action bins — turning the actor update into a cross-entropy-like objective — combined with a regularized residual network, yields better and more stable on-policy learning.

Key Contributions

1. Actor parameterization revisited as a first-class design choice. The paper studies discretized categorical policies for on-policy continuous control, where each action dimension is represented as a distribution over discrete bins so that the policy loss resembles classification cross-entropy, and analyzes the resulting policy-gradient variance.

2. A regularized actor network for discrete on-policy learning. The authors pair the categorical actor with a residual feedforward architecture using pre-layer normalization, an inverted-bottleneck feedforward block, and a linear output projection producing d × K logits, with the aim of improving optimization stability and reducing gradient variance.

3. Broad empirical demonstration across algorithms and benchmarks. RN-D is shown to improve performance and accelerate convergence on Gym locomotion and ManiSkill (both state-based and RGB-based observation), and the benefit carries over to PPO, TRPO, and SPO, indicating the gain is not tied to a single on-policy surrogate objective.

4. Ablations isolating why the method works. Loss-swap, architecture-component, network-width, bin-count, and runtime studies are used to argue that the cross-entropy-like objective, not discretization alone, drives most of the improvement.

Main Findings

• RN-D is best across benchmarks. On Gym locomotion, RN-D attains higher TD3-normalized returns than the Gaussian baselines, with clear improvements on harder tasks such as Ant and Humanoid. On ManiSkill state-based tasks, RN-D yields the highest average success across the 10 tasks, and is the only variant that reliably solves the challenging StackCube task, where other methods remain near-zero success for most of training. The same trend holds for the 5 vision-based ManiSkill tasks with CNN encoders.

• Fewer steps for the same performance. The aggregate learning curves indicate a sample-efficiency speedup: RN-D matches the best baseline's performance using roughly 1.3–1.9× fewer interaction steps.

• Ablations show two complementary effects. Comparing RN-C versus MLP-C shows the regularized backbone improves performance even with Gaussian actors; comparing MLP-D versus MLP-C shows discretization alone can also help. For vision-based tasks, the CNN encoder narrows the gap among actor variants, and MLP-C is competitive in some cases, but RN-D still achieves the best final performance and consistently better sample efficiency.

• Gradient variance is consistently lower for categorical actors. On five MuJoCo tasks, the Gaussian policy (RN-C/MLP-C) exhibits consistently larger policy-gradient variance, often by orders of magnitude, than the discrete counterparts (RN-D/MLP-D). The disparity is most pronounced for the standard MLP actor: MLP-C stays far above MLP-D and increases steadily with environment steps, consistent with the Gaussian standard deviation shrinking over training. RN-D achieves the lowest policy-gradient variance throughout training.

• Theoretical variance comparison. Proposition 4.1 shows that for a one-step REINFORCE estimator with fixed state and constant return R, the expected squared gradient norm for a Gaussian is R² Tr(Σ⁻¹) = Σᵢ R²/σᵢ², while for the factorized categorical policy it is R² Σᵢ (1 − ‖pᵢ(s)‖²) ≤ mR²(1 − 1/K), with the bound tight only when each pᵢ(s) is uniform over the K bins. The paper notes that even when the Gaussian standard deviation is relatively large, e.g. σ = 1, the per-dimension gradient-variance gap is lower bounded by 1/K and the total gap scales linearly with action dimension m.

• Loss-swap ablation identifies the objective as the key ingredient. When the discrete actor's logits are instead converted into an expected action (probability-weighted average of bin centers) and trained with Gaussian log-likelihood (a weighted squared-error objective), most of the performance improvement disappears. Discretizing the action space alone does not explain the gains.

• Component analysis links gradient quality to return. On Gym locomotion, residual connections and layer normalization each improve the gradient signal-to-noise ratio (mean SNR over mini-batches) relative to a plain MLP actor, and higher SNR correlates with higher 95th-percentile normalized return, with RN-D achieving the best of both.

• Instance-level results are presented graphically. Exact per-task normalized returns and success rates are shown in figures across seeds rather than as numeric tables in the paper content; the reported aggregate values are qualitative descriptions (higher/lower, near-zero, near-saturated).

Methodology in Plain English

The authors keep the RL algorithm and the critic network fixed and change only the actor. Instead of outputting the mean and standard deviation of a Gaussian, the actor outputs a set of logits for each action dimension, with the action range [-1, 1] split uniformly into K bins. During training, the agent samples a bin per dimension from an independent softmax, and the action taken is the corresponding bin center. Because each dimension is a separate categorical distribution, the policy's log-likelihood becomes a categorical probability, so the PPO-style update becomes a weighted cross-entropy over the selected bins — the weights coming from the advantage times the clipped importance ratio — rather than the weighted squared error a Gaussian likelihood produces.

They then replace the shallow MLP actor with a regularized residual network: an encoder (MLP for low-dimensional state, CNN for images) followed by a stack of residual feedforward blocks that use pre-layer normalization and an inverted-bottleneck design (hidden width expanded from d_h to 4d_h, ReLU, then projected back), and finally a linear layer producing d × K logits. Everything else — rollout collection, advantage estimation, the clipped surrogate objective, and the critic architecture — is held identical. They compare four variants under PPO: RN-D (regularized network + discrete actor), RN-C (regularized network + Gaussian actor), MLP-C (standard MLP + Gaussian actor, the usual default), and MLP-D (standard MLP + discrete actor), using K = 41 bins in the main experiments.

Why This Matters

Impact on research. The paper reframes actor design as an optimization problem that can borrow directly from supervised learning, where cross-entropy is a well-established drop-in for MSE and where architectural regularization is known to matter. If discretized categorical actors reliably reduce gradient variance and scale better with network capacity, that suggests a rethinking of the default Gaussian actor assumption across on-policy RL, and it connects reinforcement learning from human feedback (where policies are already categorical over tokens) with continuous-control actor–critic design.

Real-world applications (each follows from continuous-control policies that learn from interaction):

• Robotic manipulation, including the dexterous tasks in the ManiSkill suite such as the StackCube task the paper highlights.
• Legged and humanoid locomotion controllers, the setting of the Gym locomotion benchmarks used here.
• Vision-based robot control, where policies consume raw RGB inputs through a convolutional encoder.
• Any on-policy training pipeline that must be stable and sample-efficient, such as simulation-to-real control, where fewer interaction steps directly translates into lower cost.

Industry relevance. The change is described as a drop-in replacement requiring no new algorithm, and the runtime table shows MLP-D and RN-D achieve comparable or slightly better throughput than their continuous counterparts with similar wall-clock time (MLP-C: 1.65 h, 854.4 samples per second; MLP-D: 1.56 h, 925.3 SPS; RN-C: 1.75 h, 811.1 SPS; RN-D: 1.64 h, 873.8 SPS, averaged over 5 Gym locomotion tasks with 5 seeds each). The authors only claim contextualized wall-clock efficiency against TD3 and TD-MPC2 on Humanoid-v4 and StackCube-v1, not state-of-the-art among all RL algorithms — TD-MPC2 is described as considerably less computationally efficient in their setting, and TD3 performs well on Humanoid-v4 but struggles on StackCube-v1.

Future Directions

• Extend beyond PPO. The paper demonstrates gains under PPO, TRPO, and SPO, but calls for evaluating discretized categorical actors under a broader set of on-policy algorithms.
• Study actor–critic interactions. The current work deliberately isolates actor-side effects by keeping the critic fixed; jointly scaling or regularizing both networks may reveal more symmetric design principles for stable on-policy optimization.
• Move to harder embodied settings. Long-horizon manipulation, contact-rich control, and multi-object tasks are named as regimes where the benefits and limitations of the approach are not yet characterized.
• Resolve the bin-granularity question. The bin-count ablation varies K from 11 to 1001 (with 41 used in main experiments) and finds moderate counts generally best with degradation at very large counts, and that the effective range can be task-dependent; the paper notes performance is often stable around 11–101 bins, but the optimal choice varies across environments.

Target Audience

Reinforcement learning researchers and practitioners who work on continuous control, especially those using PPO-style on-policy methods and looking for simple, architecture-level improvements that do not require changing the algorithm. It is also relevant to engineers deploying on-policy RL for robotics or locomotion, and to readers interested in the connection between classification-style objectives and policy optimization, since the paper's core argument is that the likelihood model defining the actor's log-probability shapes the optimization behavior more than is commonly assumed.

Authors’ abstract

On-policy Reinforcement Learning (RL) remains a dominant paradigm for continuous control, yet standard implementations rely on Gaussian actors and relatively shallow MLP policies, often leading to brittle optimization when gradients are noisy, and policy updates must be conservative. In this paper, we revisit actor policy representation as a first-class design choice for on-policy RL. We study discretized categorical actors, which represent each action dimension as a distribution over discrete bins and induce a policy objective analogous to classification cross-entropy loss. Building on architectural advances from supervised learning, we further pair discretized categorical actors with regularized networks, yielding RN-D. Across diverse continuous-control benchmarks, we show that simply replacing the standard Gaussian actor with our proposed actor substantially improves performance, achieving state-of-the-art results within on-policy RL. We release our code at https://github.com/alwaysbyx/RND-RL.