EMCES: How Episodic Memory Guides Controllable Sample Synthesis to Boost Reinforcement Learning

Motivation

Reinforcement learning (RL) achieves impressive results in games, embodied intelligence, and large language models, yet acquiring high‑quality samples in real‑world settings remains costly and risky. Sample augmentation, especially diffusion‑based synthesis (e.g., SynthER), can expand training data but does not guarantee that the generated samples are most beneficial for policy learning.

EMCES Overview

EMCES (Episodic Memory‑Guided Controllable Experience Synthesis) introduces three key components:

A controllable diffusion model conditioned on episodic‑memory‑derived signals.

An episodic‑memory‑based temporal‑difference error (EMTD) that prioritizes samples with high potential for policy improvement.

A hash‑based state representation that makes the memory both compact and fast to query.

1.1 Controllable Diffusion Model with Episodic Memory

The model treats each RL transition (state, action, reward, next state) as a data unit. It learns a conditional diffusion process p_θ(x_t|x_{t-1},c) where the condition c encodes high‑value information extracted from episodic memory. A compact state encoder φ(s) reduces redundancy, especially in high‑dimensional visual states, and its output becomes part of the condition.

To capture richer context, the condition also incorporates a state‑action value estimate Q(s,a) or the discounted return G(s), providing the diffusion model with guidance about future rewards.

1.2 EMTD‑Based Prioritized Condition Sampling

EMTD is defined as the temporal‑difference error between the discounted return estimated from the next state and the historical optimal return stored in memory. Formally,

Larger EMTD values indicate that a sample can potentially yield higher returns than existing experience, so the sampling probability for each condition is computed by applying a softmax to EMTD, scaled by a temperature parameter β. When β=0, the strategy reduces to uniform sampling, preserving diversity.

1.3 Hash‑Based State Representation for Episodic Memory

To make the memory efficient, EMCES learns a data‑dependent binary hash for each state using the IsoHash method. A projection function h_i(s) maps the original high‑dimensional state to a single bit; concatenating all bits yields a compact code. This representation drastically reduces storage (≈8000×) and retrieval time (≈25.5×) compared with prior state encodings while retaining discriminative power.

The memory is implemented with a KD‑tree, whose storage, query, and build complexities depend on the number of encoded states N, the bits per dimension b, and the code length d. The hash‑based scheme achieves lower complexities across all three metrics.

Experimental Evaluation

Offline RL

Using D4RL benchmarks (HalfCheetah, Walker2d, Hopper, Maze2D) and three offline algorithms (TD3+BC, IQL, EDAC), EMCES‑augmented datasets consistently improve normalized scores, often surpassing the original dataset performance.

Online RL

In six online environments (quadruped‑walk, reacher‑hard, cheetah‑run, Walker2d, HalfCheetah, Hopper) with SAC as the base algorithm, EMCES outperforms both SynthER and the online‑focused PGR method, demonstrating higher sample efficiency and faster convergence.

Ablations

Comparisons of different state representations show that the hash‑based encoding retains downstream performance while reducing memory usage by ~8000× and time overhead by ~25.5×. Additional ablations confirm the importance of the episodic‑memory‑driven condition design and the EMTD‑based sampling strategy.

Conclusion

EMCES provides a strongly controllable sample synthesis pipeline that generates higher‑quality RL experiences, leverages episodic memory to prioritize valuable samples, and employs a hash‑based state representation to achieve substantial efficiency gains, making it a practical tool for both offline and online reinforcement‑learning scenarios.