Probabilistic Embeddings for Actor-Critic RL

Infer which task you face from a few recent steps, then run one task aware policy from that belief with no fresh gradients.

Three rewards 0.8, 1.0, 0.6 compress to belief N(0.90, 0.25), and the conditioned policy steps plus 0.72 toward the right goal

Why Does This Exist?

Fresh training per task is ruinous. Move a cheetah goal velocity from $1.0$ to $2.5$ meters per second and a fixed policy fails, so standard deep RL burns millions of new steps to relearn locomotion it already knew. Real robots and heavy simulators cannot afford that per variant.

Meta-Reinforcement Learning trains across a task family so a new member costs episodes, not millions of steps. Early answers paid their own tax: MAML in RL needs fresh on policy rollouts per meta update, and recurrent black box adapters tangle the adaptation rule inside memory that reuses poorly off policy.

PEARL (Rakelly and coauthors, ICML 2019) takes the inference view. Train one off policy Soft Actor-Critic agent plus a context encoder that reads recent transitions and infers a compact task variable. Meta-training reuses replay data like SAC does; meta-testing adapts by sharpening a small belief, not by differentiating through the policy.

Think of It Like This

The substitute chef who tastes first

You cover a kitchen with two regulars: one orders extra salt, one orders none. Nobody tells you who sits at table six. You do not reread the cookbook. You taste three spoonfuls: fairly salty, very salty, mildly salty. You average the evidence onto one card reading salty strength about $0.9$ of $1$ , clip it to your station, and cook the whole service from it.

Tomorrow the spoonfuls taste flat. Same averaging, opposite card, same hands. You commit to one card per service instead of second guessing after every bite.

Where it stops: no palate is a permutation invariant net and no chef pays a KL penalty to stay near a prior. The commit per service carries over; that is posterior sampling for one episode.

How It Actually Works

A task $\mathcal{T}$ is one MDP from a family $p(\mathcal{T})$ . Plain words: the family is the menu of kitchens, tonight is one kitchen. Context $c$ is a small bag of transitions $(s, a, r, s^{\prime})$ with state $s$ , action $a$ , reward $r$ , next state $s^{\prime}$ . Latent variable $z$ is a short vector naming the task. Encoder $q_{\phi}(z \mid c)$ with weights $\phi$ maps context to a distribution over $z$ . Policy $\pi_{\theta}(a \mid s, z)$ with weights $\theta$ maps state plus task sample to actions. Critic $Q(s, a, z)$ scores state, action, and task guess together.

Context encoder and conditioned SAC

Each transition passes through a shared net $f_{\phi}$ giving one Gaussian factor. The encoder multiplies those factors with a standard normal prior $\mathcal{N}(0, I)$ and renormalizes. Actor and critic are plain SAC nets that concatenate $z$ to their inputs. Training draws two batches per task from separate buffers: a context batch for $z$ and an RL batch for Bellman updates. The critic minimizes soft Bellman error plus $\beta \, \mathrm{KL}(q_{\phi}(z \mid c) \parallel \mathcal{N}(0, I))$ , where $\beta$ caps carried task detail and $\mathrm{KL}$ measures distance from the prior. The actor maximizes the SAC objective with entropy temperature $\alpha$ . The encoder learns from the critic loss only; actor gradients proved too noisy in the paper ablations. Testing samples one $z$ per episode and holds it fixed, which is Thompson sampling: broad prior tries something coherent, growing context sharpens the posterior into task locking behavior.

Worked example

Line world, goals at $+1$ (task R) and $-1$ (task L). At $s = 0.0$ with three rightward probes $c_1 = (0.0, +0.3, 0.8)$ , $c_2 = (0.3, +0.3, 1.0)$ , $c_3 = (0.6, +0.2, 0.6)$ , toy factor means $m_i = 1.5 r_i$ give $m = [1.20, 1.50, 0.90]$ summing to $3.60$ . Prior precision $1$ plus three unit factors gives total precision $4$ , variance $1 / 4 = 0.25$ , deviation $0.50$ , mean $3.60 / 4 = 0.90$ . So $q(z \mid c) = \mathcal{N}(0.90, 0.25)$ . Sample $z = 0.90$ ; toy policy mean $\mu_a = 0.8 z + 0.1 s$ at $s = 0.0$ gives $\mu_a = 0.72$ toward $+1$ . The mirror rewards $[-0.8, -1.0, -0.6]$ give mean $-0.90$ and action $-0.72$ . The diagram draws this exact chain. Belief cost against the prior is about $0.72$ nats before $\beta$ weighting.

Code

The snippet replays both task directions with the standard library only.

1def encode_context(rewards: list[float], scale: float = 1.5) -> tuple[float, float]:2    """Product of unit variance factors plus standard normal prior."""3    total = sum(scale * r for r in rewards)4    precision = 1.0 + len(rewards)5    return total / precision, 1.0 / precision6
7
8def policy_mean(z: float, s: float = 0.0) -> float:9    return 0.8 * z + 0.1 * s10
11
12mu, var = encode_context([0.8, 1.0, 0.6])13print(round(mu, 2), round(var, 2), round(policy_mean(mu), 2))  # -> 0.9 0.25 0.7214
15mul, varl = encode_context([-0.8, -1.0, -0.6])16print(round(mul, 2), round(varl, 2), round(policy_mean(mul), 2))  # -> -0.9 0.25 -0.72

Both prints match the hand derived posteriors $\mathcal{N}(0.90, 0.25)$ and $\mathcal{N}(-0.90, 0.25)$ with actions $+0.72$ and $-0.72$ .

Watch Out For

Testing on a context the encoder never saw

Training curves look healthy, then new tasks adapt slowly or pick the wrong side for episodes. Usual cause is context shift: training sampled large uniform batches of hundreds of transitions while testing feeds the last five online steps. The encoder meets a new size and distribution, its mean lands between tasks, variance stays wide.

Match the pipeline: same context batch sizes at train and test, recent plus uniform mixes as in the paper sampler, context batch separate from the RL batch. Log posterior mean and variance on held out tasks; test variance far above training variance means the pipeline drifted.

z-collapse: the policy stops listening to z

Posterior means sit near zero for every task, variances hug the prior, actions ignore left versus right. Either bottleneck $\beta$ prices any nonzero $z$ out of existence, or encoder gradients flow through the actor loss and drown the task signal in policy noise.

Train the encoder through the critic loss only, start $\beta$ small and raise until task clusters just hold, and plot $z$ means colored by true task. Distinct clusters carry signal; one blob means the bottleneck is shut and lowering $\beta$ comes before resizing nets.

The Quick Version

PEARL pairs one SAC agent with a context encoder so a single policy serves a task family via latent $z$ .
The encoder turns a bag of transitions into Gaussian belief $q(z \mid c)$ , sampled once per episode for coherent search.
Separate context and RL batches from replay keep sample cost far below on policy meta-RL.
Rewards $[0.8, 1.0, 0.6]$ give $\mathcal{N}(0.90, 0.25)$ and mean action $+0.72$ toward the right goal.
Context shift and $z$ -collapse are the first two checks when adaptation stalls or ignores the task.