Research map / Reinforcement learning / Policy gradient and actor-critic
Time-Efficient Reinforcement Learning with Stochastic Stateful Policies
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Stochastic Stateful Policy Gradient (S2PG), which trains stateful policies by modeling a stochastic internal-state kernel jointly with a stateless policy, avoiding BPTT, with theory and MuJoCo/locomotion experiments.
Problem
Training stateful (recurrent) RL policies with BPTT is slow, sequential, suffers from vanishing/exploding gradients, and requires truncation that biases the gradient.
Approach
Decompose a stateful policy into a stochastic internal state kernel and a stateless policy that jointly output actions and next internal states, and derive stateful policy gradient theorems so gradients depend only on local transitions (no backpropagation through time), enabling stateful versions of SAC, TD3, PPO, GAIL, and LSIQ.
Key results
- S2PG-based agents trade a slight drop in asymptotic performance for SAC and TD3 but train much faster than BPTT with long histories, performing better on high-dimensional tasks like Ant and Humanoid.
- On imitation learning locomotion tasks (Atlas, Humanoid with dynamics randomization), S2PG outperforms BPTT baselines even in terms of samples.
- Without privileged information, the PPO-based S2PG variant reliably outperforms BPTT and window approaches and is faster than SLAC.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace BPTT with S2PG for recurrent SAC and TD3 | SAC and TD3 on privileged-critic MuJoCo tasks · MuJoCo POMDP tasks | asymptotic return | Hurt |
| replace BPTT with S2PG | Recurrent SAC and TD3 · MuJoCo POMDP tasks | training time | Helped |
| use S2PG instead of BPTT | Recurrent SAC and TD3 · MuJoCo POMDP tasks | return on high-dimensional observation spaces | Helped |
| use S2PG instead of BPTT | Recurrent SAC and TD3 · MuJoCo POMDP tasks | return on low-dimensional observation spaces | Hurt |
| use S2PG instead of BPTT | Recurrent PPO · MuJoCo POMDP tasks without privileged critic information | performance | Helped |
| use longer BPTT truncation histories | Recurrent PPO · MuJoCo POMDP tasks without privileged critic information | training time | Hurt |
| use S2PG instead of SLAC | Recurrent policy learning · MuJoCo POMDP tasks without privileged critic information | performance | Helped |
| use S2PG instead of BPTT | Stateful GAIL and LSIQ · Atlas and Humanoid locomotion imitation tasks | sample efficiency | Helped |
| use S2PG instead of BPTT | Stateful policies · Point-mass memory task | performance | Helped |
| use S2PG instead of BPTT | Stateful policies · Maze-door memory task | performance | Hurt |
| use S2PG instead of BPTT | Stateful policies · Memory tasks | training time | Helped |
Surprising findings
- On the second memory task (maze doors), S2PG is slightly weaker than BPTT, attributed to extra variance causing more exploration into absorbing states.
Limitations
- Cannot yet learn a bootstrapped value function and internal policy state simultaneously in POMDPs; off-policy methods need a privileged critic
- Higher variance than BPTT in some settings due to stochastic internal state kernel
- Slight drop in asymptotic performance for SAC and TD3 on low-dimensional tasks
Experimental setup
Datasets: MuJoCo Gym locomotion tasks LocoMujoco benchmark (Atlas and Humanoid imitation tasks) memory tasks (point mass, maze doors)
Benchmarks: MuJoCo Gym POMDP tasks LocoMujoco Atlas and Humanoid imitation tasks memory tasks
Metrics: normalized return training time for 1M steps gradient variance bounds
Baselines: BPTT with truncation length 5 BPTT with truncation length 32 SLAC stateless window of 5 observations stateless window of 32 observations vanilla (observation-only) agent oracle (full-state) agent
Builds on: BPTT REINFORCE PPO SAC TD3
Closest papers by meaning
- State Regularized Policy Optimization on Data with Dynamics Shift (NeurIPS 2023 · similarity 0.88)
- Simplified Temporal Consistency Reinforcement Learning (ICML 2023 · similarity 0.88)
- Absolute State-wise Constrained Policy Optimization: High-Probability State-wise Constraints Satisfaction (NeurIPS 2026 · similarity 0.88)
- Learning Optimal Deterministic Policies with Stochastic Policy Gradients (ICML 2024 · similarity 0.87)
- Convergence Analysis of Policy Gradient Methods with Dynamic Stochasticity (ICML 2025 · similarity 0.86)
- GenPO: Generative Diffusion Models Meet On-Policy Reinforcement Learning (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.