Symmetry Learning for Function Approximation in Reinforcement Learning

Overview

The paper proposes a framework for discovering symmetries from an agent's experience and using them to improve the sample efficiency of reinforcement learning with function approximation. A symmetry maps a state-action pair $(s,a)$ to another pair $(f(s),g_s(a))$ while preserving rewards and transition behavior. Consequently, symmetric pairs have the same role in decision making and should have the same optimal action value.

The method does not learn an explicit transformation such as a reflection or rotation. Instead, it identifies individual state-action pairs that appear equivalent because they generate similar future reward sequences. These estimated pairs are then incorporated into Q-learning through an auxiliary loss that encourages their predicted Q-values to agree. The framework therefore consists of two components: detecting candidate symmetries from reward histories and using the detected pairs as a soft prior when training the Q-function.

Methodology

Symmetry Detection

The main assumption is that symmetric state-action pairs generate the same distribution over future reward sequences when followed by corresponding policy behavior. For each state-action pair $(s,a)$ and sequence length $j$, the method records

$$\Pi_{sa,j} = \left\{ (\sigma,N_{\sigma}) \right\},$$

where $\sigma=(r_0,\ldots,r_{j-1})$ is a reward sequence observed after beginning from $(s,a)$ and $N_{\sigma}$ is the number of times that sequence was observed.

The total number of reward trails of length $j$ observed after $(s,a)$ is

$$|\Pi_{sa,j}| = \sum_{|\sigma|=j} N_{\sigma}.$$

For two state-action pairs $(s,a)$ and $(\bar{s},\bar{a})$, the intersection of their reward-trail collections is defined as

$$\Pi_{sa,j}\cap\Pi_{\bar{s}\bar{a},j} = \left\{ \left( \sigma, \min\left(N_{\sigma},\bar{N}_{\sigma}\right) \right) \right\},$$

where $N_{\sigma}$ and $\bar{N}_{\sigma}$ are the occurrence counts of reward sequence $\sigma$ for $(s,a)$ and $(\bar{s},\bar{a})$, respectively. Therefore,

$$|\Pi_{sa,j}\cap\Pi_{\bar{s}\bar{a},j}| = \sum_{|\sigma|=j} \min\left(N_{\sigma},\bar{N}_{\sigma}\right).$$

The similarity between the two state-action pairs is calculated as

$$\chi_{i,l_0} \left( \langle s,a\rangle, \langle\bar{s},\bar{a}\rangle \right) = \frac{ \displaystyle \sum_{j=l_0}^{i} \left| \Pi_{sa,j} \cap \Pi_{\bar{s}\bar{a},j} \right| }{ \displaystyle \left( \sum_{j=l_0}^{i} |\Pi_{sa,j}| \; \sum_{j=l_0}^{i} |\Pi_{\bar{s}\bar{a},j}| \right)^{1/2} }.$$

The numerator measures the overlap between the reward-trail counts of the two pairs over sequence lengths from $l_0$ to $i$. The denominator normalizes this overlap according to the total number of trails observed from each pair.

Pairs whose similarity exceeds a threshold $\Delta$ are treated as symmetric:

$$\chi_{\mathrm{sym}} = \left\{ \left( \langle s,a\rangle, \langle\bar{s},\bar{a}\rangle \right) \;\middle|\; \chi_{i,l_0} \left( \langle s,a\rangle, \langle\bar{s},\bar{a}\rangle \right) \geq \Delta \right\}.$$

The paper proves a completeness result. If $(s,a)$ and $(\bar{s},\bar{a})$ are equivalent under a true symmetry, then under the assumptions of the theorem,

$$\lim_{|\mathrm{episodes}|\rightarrow\infty} \chi_{i,l_0} \left( \langle s,a\rangle, \langle\bar{s},\bar{a}\rangle \right) = 1.$$

This establishes that true symmetric pairs can eventually be detected, but it does not establish the converse. The issue is that unrelated pairs can also receive a high similarity score if they happen to generate the same reward sequences. This is especially true for sparse rewards, where a reward sequence of zeros is experienced by many unrelated state-actions.

Potential-based reward shaping is used to make the reward histories more informative:

$$R'(s,a,s') = R(s,a,s') + F(s,a,s'),$$

where the shaping term is

$$F(s,a,s') = \gamma\Theta(s') - \Theta(s).$$

Symmetry Inclusion Priors and Learning

After identifying candidate symmetric pairs, the goal is to make the Q-function assign them identical values:

$$Q(s,a;\theta) = Q(\bar{s},\bar{a};\theta).$$

The ordinary temporal-difference loss is

$$L_{i,\mathrm{TD}}(\theta_i) = \mathbb{E}_{(s,a,r,s')\sim\mathbb{B}} \left[ \left( \left[ r + \gamma \max_{a'} Q(s',a';\theta_{i-1}) \right] - Q(s,a;\theta_i) \right)^2 \right].$$

The paper introduces a symmetry loss

$$L_{i,\mathrm{Sym}}(\theta_i) = \mathbb{E}_{ \left( \langle s,a\rangle, \langle\bar{s},\bar{a}\rangle \right) \sim \chi_{\mathrm{sym}} } \left[ \left( Q(\bar{s},\bar{a};\theta_i) - Q(s,a;\theta_i) \right)^2 \right].$$

The complete training objective is

$$L_{i,\mathrm{total}}(\theta_i) = L_{i,\mathrm{TD}}(\theta_i) + \lambda L_{i,\mathrm{Sym}}(\theta_i),$$

where $\lambda$ controls the strength of the symmetry prior.

Limitations

• The detector compares reward sequences using exact equality. It is therefore most suitable for deterministic or stochastic rewards with a discrete set of values. Independently sampled continuous rewards will almost never match exactly without quantization or other measures.
• Reward trails depend on the policy used after the initial state-action pair. The theoretical argument relies on corresponding behavior in symmetric regions, while an arbitrary policy early in training may generate different trails from genuinely symmetric pairs.
• The method stores and compares histories for individual state-action pairs. This works for small discrete spaces, but is impractical for large and continuous spaces.