Technical Summary: Simplicial Embeddings Improve Sample Efficiency in Actor–Critic Agents

Problem Statement

Deep reinforcement learning (RL) faces a persistent tension between wall-clock training speed and sample efficiency. While recent advancements in massive environment parallelization (e.g., FastTD3, SEED RL) have accelerated training time, these methods often require a prohibitively large number of environment interactions to achieve desired performance levels. This inefficiency is particularly problematic in domains where interactions are expensive (e.g., real-world robotics) or where data-hungry algorithms fail to transfer from simulation to reality.

A core underlying issue identified in the paper is representation collapse exacerbated by non-stationarity. In actor–critic methods, the critic is trained against bootstrapped targets that shift as the policy updates. This non-stationary training signal can cause learned representations to lose diversity, leading to "neuron dormancy," ill-conditioned regression problems for the critic, and unstable value estimates. Standard deep RL architectures often lack geometric constraints to prevent these pathologies, resulting in poor sample efficiency and divergence.

Methodology: Simplicial Embeddings (SEM)

The authors propose Simplicial Embeddings (SEM), a lightweight architectural component designed to impose a geometric inductive bias on the latent representations of actor–critic networks.

• Mechanism: SEM partitions the encoder's output vector into $L$ groups of size $V$. Within each group, a softmax operation is applied to the logits, constrained by a temperature parameter $\tau$.
  $$\tilde{z}_{\ell,v} = \frac{\exp(z_{\ell,v}/\tau)}{\sum_{v'=1}^V \exp(z_{\ell,v'}/\tau)}$$
• Geometric Constraint: This transformation forces the latent vector to lie on a product of simplices ($\Delta^{V-1} \times \dots \times \Delta^{V-1}$).
• Properties:
  1. Boundedness: Group-wise normalization ensures embeddings remain bounded, preventing the extrapolation of arbitrarily large Q-values that destabilize critics.
  2. Sparsity: The softmax competition (especially at low $\tau$) induces sparse, near-one-hot encodings, promoting discrete-like features without the bias of straight-through estimators.
  3. Differentiability: Unlike explicit discretization methods (e.g., VQ, Gumbel-Softmax), SEM is fully differentiable, avoiding biased gradient estimators.
  4. No Auxiliary Losses: SEM acts as a structural layer (similar to an activation function) rather than requiring auxiliary reconstruction or contrastive losses, adding negligible computational overhead.

The authors integrate SEM into the penultimate layer of both the actor and critic networks in various algorithms, including FastTD3, FastSAC, and PPO.

Key Contributions

1. Identification of Non-Stationarity as a Driver of Collapse: The paper provides empirical evidence (via CIFAR-10 experiments with shuffled labels) and theoretical analysis showing that non-stationary targets in deep RL lead to increased neuron dormancy and reduced effective rank, directly linking these phenomena to training instability.
2. Introduction of SEM for RL: The authors adapt Simplicial Embeddings from self-supervised learning to the RL domain, demonstrating that constraining representations to a product of simplices stabilizes bootstrapping and strengthens policy gradients.
3. Comprehensive Empirical Validation: The method is evaluated across a diverse suite of benchmarks, including:
   • Continuous Control: HumanoidBench (28 tasks), Isaac Gym, and MTBench.
   • Discrete Control: Extended Atari-10 (28 ALE games).
   • Algorithms: FastTD3, FastSAC, PPO, and Flow Q-Learning (FQL).
4. Analysis of Representation Dynamics: The study connects performance gains to specific representational improvements, such as higher effective rank, bounded feature norms, reduced critic disagreement, and lower TD error.

Results

• Sample Efficiency: Across all tested algorithms and environments, agents equipped with SEM consistently achieve higher returns with fewer environment steps compared to baselines. For instance, in HumanoidBench, SEM accelerates early learning and improves asymptotic performance.
• Robustness: SEM improves performance even when core design choices of fast RL algorithms are altered (e.g., removing distributional critics, reducing batch sizes, or using smaller replay buffers).
• Stability: SEM agents exhibit lower variance across random seeds, smaller TD errors, and more calibrated value estimates. The method effectively mitigates the "moving target" problem in actor–critic training.
• Parameter Sensitivity: The total representational capacity ($L \times V$) is the primary driver of performance. In high-capacity regimes, performance saturates, but SEM remains effective even with smaller dimensions (e.g., $V=4$) in certain contexts.
• Comparison: SEM outperforms alternative methods for inducing structure, such as Gumbel-Softmax with straight-through estimation and Vector Quantization, likely due to the avoidance of biased gradient estimators.

Significance and Claims

The paper claims that representation geometry is a critical, yet often overlooked, lever for stabilizing deep RL under non-stationarity. By enforcing a simplex-based structure, SEM provides a "simple but powerful" mechanism to prevent feature collapse and neuron dormancy without the computational cost of auxiliary losses or model-based world models.

The authors emphasize that:

• SEM is computationally efficient, adding effectively zero runtime cost while significantly improving sample efficiency.
• The benefits are most pronounced when applied to the actor's penultimate layer, where feature geometry directly shapes policy gradients, though applying it to the critic also yields gains.
• This approach offers a path toward more robust and data-efficient RL agents, particularly in scenarios where interactions are costly or distribution shifts are inevitable.

The paper concludes that while SEM is not a universal remedy for all RL challenges (e.g., extreme sparsity or massive distribution shifts may still require tuning), it represents a significant step toward inducing order "from within" the representation space, aligning with the philosophical stance that stability emerges from structural constraints rather than solely from increased compute or heavier regularization.