Simplicial Embeddings Improve Sample Efficiency in Actor-Critic Agents

TL;DR

SEM constrains encoder outputs to a product of simplices $(\Delta^{V-1})^L$, providing a sparse, bounded geometric prior that stabilizes bootstrapped targets under non-stationarity. Applied to FastTD3, FastSAC, and PPO, SEM improves data efficiency and final performance across continuous and discrete tasks with negligible overhead. The authors show SEM preserves actor/critic representation diversity (increased effective rank, bounded feature norms) and reduces TD error and critic disagreement, mitigating representation collapse. This work highlights representation geometry as a practical, drag-free regularizer that enhances robustness and transferability in RL.

Abstract

Recent works have proposed accelerating the wall-clock training time of actor-critic methods via the use of large-scale environment parallelization; unfortunately, these can sometimes still require large number of environment interactions to achieve a desired level of performance. Noting that well-structured representations can improve the generalization and sample efficiency of deep reinforcement learning (RL) agents, we propose the use of simplicial embeddings: lightweight representation layers that constrain embeddings to simplicial structures. This geometric inductive bias results in sparse and discrete features that stabilize critic bootstrapping and strengthen policy gradients. When applied to FastTD3, FastSAC, and PPO, simplicial embeddings consistently improve sample efficiency and final performance across a variety of continuous- and discrete-control environments, without any loss in runtime speed.

Figures (21)

• Figure 1: Training dynamics on CIFAR-10 with stationary vs. non-stationary targets. In the stationary regime (fixed targets), losses decrease smoothly, neuron dormancy and effective rank remains controlled, suggesting stable representation learning. In the non-stationary regime (targets shuffled every 20 epochs), the model exhibits higher variance in losses, increased dormant neuron rates, and reduced effective rank. The addition of SEM mitigates this instability.
• Figure 2: Actor–critic network architecture with SEM. The actor (left) and critic (middle) architectures are modified with a SEM module, which partitions features into groups and applies group-wise softmax (right panel), constraining them to a product of simplices.
• Figure 3: Average normalized return on $5$ HumanoidBench tasks over $6$ seeds. Baseline agent (blue, - -) vs. SEM variants applied to actor, critic, or both. Each curve corresponds to an embedding dimension; dim$=64$ (green, ---) is highlighted. SEM accelerates early learning and improves asymptotic performance, with dim$=64$ giving the most stable gains.
• Figure 4: Learning and representation diagnostics on $2$ HumanoidBench tasks over $6$ seeds. SEM reaches high return earlier, raises actor/critic effective rank, and keeps actor features compact.
• Figure 5: Learning dynamics on $2$ HumanoidBench tasks. SEM reaches high return faster, with lower losses, smaller TD error, reduced critic disagreement, and better-calibrated value estimates.

Theorems & Definitions (2)

• Theorem 1
• proof