Abstraction for Offline Goal-Conditioned Reinforcement Learning

Offline RL is commonly formulated as a two-step procedure: first, a value function is learned, and, second, a policy is extracted from this value function. A major challenge in offline RL is extracting an optimal policy from regions of the state-space where the dataset only contains low-quality transitions. This is because, to prevent action value overestimation, model-free algorithms typically either incorporate some form of conservatism during value-learning, or regularise policy extraction towards the dataset distribution using behaviour cloning.

Recent Rejection-Sampling-based methods take this to the extreme: by learning the dataset-generating policy using behaviour cloning, and then, at deployment, sampling actions from this policy and choosing those with the highest value, the deployed policy will never take higher-quality actions than those contained within the dataset.

In principle, our aim is to learn abstractions that enable the reuse of experience across similar contexts. By learning options that group together state-waypoint pairs which, under an optimal policy, induce similar action sequences, options would be relative rather than anchored to an absolute frame of reference. This would naturally induce two hierarchical notions of similarity: a high-level similarity, where similar state-goal pairs induce similar options, and a low-level similarity, where similar state-option pairs induce similar immediate actions.

Our aim would be to learn high-level embeddings $\phi_h(s,g)$ and low-level embeddings $\phi_l(s,\omega)$ that abstract away information irrelevant to their respective levels: $$ \pi(a \mid s,g) = \pi_l(a \mid \phi_l(s, \omega)) \quad \text{where} \quad \omega \sim \pi_h(\cdot \mid \phi_h(s,g)). $$

We now use the above framework to guide algorithmic decision making.

• First, allowing the high-level and low-level decision processes to operate on different representations and at different temporal abstractions (discount factors), requires two value functions.
• Second, options should be learned via action similarity rather than, for example, value-similarity.

Abstractive RL Implicitly Learning Relativised Options (ARLi)

We make the most minimal amendment to Hierarchical Implicity Q-Learning (HIQL) to comply with the ARL framework. We learn two value functions: one for the high-level policy and one for the low-level policy. Following HIQL, the option space is bounded to a hypersphere to introduce geometric regularisation, rather than simply learning a distinct embedding for every possible state-waypoint pair. However, rather than learning option representations jointly with the value function, we learn option representations with the low-level policy. Intuitively, we want to push together state-waypoint pairs that require similar low-level actions rather than state-waypoint pairs with similar values.

The results are promising! Without any hyperparameter tuning (we simply inherit those tuned for HIQL), we achieve non-trivial performance in many of the higher-dimensional OGBench environments.

Abstractive RL Explicitly Imposing Translation Invariance (ARLe)

At a much more extreme level, we now impose translation invariance on the low-level value function to explicitly enforce learning from similar contexts across the state-space. We use relativised states, which we define as $v = g_s - s$, where $s$ is the current state and $g_s$ a waypoint to the goal, and explicit relativised options to encourage abstraction from the absolute frame of reference. We do still condition the low-level policy on the absolute state, however, to enable the policy to satisfy local constraints.

Notably, ARLe outperforms ARLi in two out of the four sparsest datasets:

The results are less impressive than ARLi (this rigid enforcement is not ideal in practice and will not scale: the representation is too lossy!), but what is the potential?

It is a proof of concept:

Why do we not tune hyperparameters for ARL?
Offline RL typically requires significant online tuning, which is expensive and would be unscalable for training a billion-parameter general-purpose agent. Scaling offline RL should be robust to hyperparameters, which is why our work focuses on inductive biases rather than representational losses that introduce additional hyperparameters.

Why does ARL still perform poorly in certain tasks?
There are many confounding factors! Some examples: poor high-level value learning and a lack of gradient in long horizon tasks; the high-level policy being unimodal rather than multimodal; not training for enough gradient steps for large dataset sizes; not learning goal representations for the high-level policy. The aforementioned issues could be mitigated by combining ARL with value horizon reduction methods such as TD-$n$ or Transitive RL, learning a Flow Policy rather than a Gaussian (especially for the high-level policy), training for more steps (we run all experiments for 1M), and using representation learning for the high-level decision process.

What environments does ARL have most impact in?
Both of our ARL-guided algorithms exhibit positive improvement trends over our benchmark algorithms with increasing dataset sparsity, which would make sense; relativised options and experience reuse become increasingly important as sparsity increases.

This work just acts as a proof of concept that relativised options and inductive biases can improve generalisation in offline GCRL.

Now: can we impose more flexible inductive biases? Or leverage action chunking to learn relativised options over action sequences?

Thank you to Seohong Park, whose work inspired many of these ideas.