Contrastive Retrospection: honing in on critical steps for rapid learning and generalization in RL
Authors: Chen Sun
Wannan Yang
Thomas Jiralerspong
Dane Malenfant
Benjamin Alsbury-Nealy
Yoshua Bengio
Blake Aaron Richards
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
long term credit assignment
rapid credit assignment
contrastive learning
few-shot learning in RL
TL;DR: ConSpec uses offline contrastive learning to hone in on the critical steps of RL tasks, helping both rapid long term credit assignment and out-of-distribution generalization.

Abstract: In real life, success is often contingent upon multiple critical steps that are distant in time from each other and from the final reward. These critical steps are challenging to identify with traditional reinforcement learning (RL) methods that rely on the Bellman equation for credit assignment. Here, we present a new RL algorithm that uses offline contrastive learning to hone in on these critical steps. This algorithm, which we call Contrastive Retrospection (ConSpec), can be added to any existing RL algorithm. ConSpec learns a set of prototypes for the critical steps in a task by a novel contrastive loss and delivers an intrinsic reward when the current state matches one of the prototypes. The prototypes in ConSpec provide two key benefits for credit assignment: (i) They enable rapid identification of all the critical steps. (ii) They do so in a readily interpretable manner, enabling out-of-distribution generalization when sensory features are altered. Distinct from other contemporary RL approaches to credit assignment, ConSpec takes advantage of the fact that it is easier to retrospectively identify the small set of steps that success is contingent upon (and ignoring other states) than it is to prospectively predict reward at every taken step. ConSpec greatly improves learning in a diverse set of RL tasks. The code is available at the link: https://github.com/sunchipsster1/ConSpec

Contrastive Retrospection: honing in on critical steps for rapid learning and generalization in RL
Chen Sun"
Wannan Yang New York University
Thomas Jiralerspong
Mila, Université de Montréal sunchipsster&gmall. com
Mila, Université de Montréal thomas. jiralerspong @mila. quebec
vinnleyangm@@@mall.com
Dane Malenfant
Benjamin Alsbury-Nealy
McGill University
University of Toronto, SilicoLabs Incorporated benj amin. alsbury nealy@si aillsilllbllss
dane malenfant@mil quubbbbeeeceeaecccc quebec
Yoshua Bengio
Blake Richards
Mila, Université de Montréal, CIFAR feshua.bengiodien quueeee
Mila, McGill University
Learning in Machines & Brains, CIFAR blake riceerrihheee chards@mila. quebec
Abstract
In real life, success is often contingent upon multiple critical steps that distant time from each other and from the final reward. These critical steps are challenging identify with traditional reinforcement learning (RL) methods that rely on the Bellman equation for credit assignment. Here, we present new RL algorithm that uses offline contrastive learning to hone on these critical steps. This algorithm, which we call Contrastive Retrospection (ConSpec), can be added to any existing RL algorithm. ConSpec learns set of prototypes for the critical steps in a task by novel contrastive loss and delivers intrinsic reward when the current state matches one the prototypes. The prototypes in ConSpec provide two key benefits for credit assignment: (i) They enable rapid identification of all the critical steps. (ii) They do so in a readily interpretable manner, enabling out-of-distribution generalization when sensory features altered. Distinct from other contemporary RL approaches to credit assignment, ConSpec takes advantage of the fact that it is casier to retrospectively identify the small set of steps that success is contingent upon (and ignoring other states) than is to prospectively predict reward at every taken step. ConSpec greatly improves learning diverse set of RL tasks. The code is available at the link: http:://gthub.com/suceppststtr/t
Introduction
In real life, succeeding in a given task often involves multiple critical steps. For example, consider the steps necessary for getting a paper accepted at conference. One must (i) generate good idea, (ii) conduct mathematical analyses or experiments, (iii) write a paper, and finally, (iv) respond to reviewers in satisfactory manner. These the critical steps necessary for success and skipping any of these steps will lead to failure. Humans are able to learn these specific critical steps even though they interspersed among many other tasks in daily life that are not directly related the goal. Humans are also able to generalize knowledge about these steps to myriad of new projects throughout an academic career that can span different topics, and sometimes, even different fields.
"Co-corresponding authors
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Though understudied in RL, tasks involving multiple contingencies separated by long periods of time capture an important capability of natural intelligence and a ubiquitous aspect of real life.
But even though such situations are common in real life, they are highly nontrivial for most RL algorithms to handle. One problem is that the traditional approach for credit assignment in RI, using the Bellman equation (58) 77. will take a lot of time to propagate value estimates back to the earlier critical steps if the length of time between critical steps is sufficiently long (48) 46). Even with the use of strategies to span a temporal gap (e.g. TD(A),, Bellman back-ups do not scale to really long-term credit problems. In light of this, some promising contemporary RL algorithms have proposed mechanisms beyond Bellman backup to alleviate the problem of long term credit assignment (3) 28 48 [2] 10). But as we show (Figs. and and|
A even these contemporary RL methods are often insufficient to solve simple instantiations of tasks with multiple contingencies, indicating that there are extra difficulties for long-term credit assignment when multiple contingencies are present.
As a remedy to this problem, we introduce Contrastive Retrospection (ConSpec), which can be added to any backbone RL algorithm. Distinct from other approaches to credit assignment, ConSpec takes advantage of the assumption exploited by humans that success is often contingent upon a small set of steps: it is then easier to retrospectively identify that small of steps than it is to prospectively predict reward at every step taken in the environment. ConSpec learns offline from a memory buffer with a novel contrastive loss that identifies invariances amongst successful episodes, i.e., the family of states corresponding to a successfully achieved critical step. To do this, ConSpec learns a state encoder and series of prototype: that represent critical steps. When the encoded representation of state is well aligned with one of the prototypes, the RI. agent receives an intrinsic reward 52 47). thereby steering the policy towards achieving the critical steps.
ConSpec is descended from a classical idea: the identification of "bottleneck states", Le. states that must be visited to obtain reward, and their use as sub-goals (38) 41 50 59 5). However, critical
steps in non-toy tasks rarely correspond to specific states of the environment. Instead, there can be large, potentially even infinite, set of states that correspond to taking a critical step in task. For example, what are the "states" associated with "conducting an experiment" in research? They cannot be enumerated but a function that detects them can be learned. Ultimately, we thus need RL agents that can learn how to identify critical steps without assuming that they correspond to specific state or easily enumerable set of states. ConSpec, by harnessing a novel contrastive loss, scalably solves this. Our contributions in this paper are as follows:
We introduce a scalable algorithm (ConSpec) for rapidly honing in on critical steps. It uses a novel contrastive loss to learn prototypes that identify invariances amongst successful episodes.
We show that ConSpec greatly improves long-term credit assignment in a wide variety of RL tasks including grid-world, Atari, and 3D environments, as well as tasks where we had not anticipated improvements, including Montezuma's 's Revenge and continuous control tasks.
We demonstrate that the invariant nature of the learned prototypes for the critical steps enable zero-shot out-of out-o--distribution generalization in RL.
2 Related work
ConSpec is a relatively simple design, but it succeeds because it centralizes several key intuitions shared with other important works. ConSpec shares with bottleneck states (38), hicrarchical RL (HRL), and options discovery (59 5 41 (50) the idea that learning to hone in sparse number of critical steps may be beneficial. But, unlike bottleneck state solutions, ConSpec does not assume that critical steps correspond to individual states (or small, finite sets of states) of the environment. In con- trast, ConSpec identifies critical steps in complex, high-dimensional tasks, such as 3D environments. As well, in HRL and options discovery, how discover appropriate sub-goals at scale remains an unsolved problem. Critical steps discovered by ConSpec could theoretically be candidate sub-goals.
ConSpec shares with Return Decomposition for Delayed Rewards (RUDDER) 3 44 62). Temporal Value Transport (TVT) (28). Synthetic Returns (SynthRs) (48), and Decision Transformers (DTs) (2) the use of Non-Bellman-base long-term credit assignment. ConSpec shares with slot-like
"We use the word "prototype" here in the psychological sense, Le. an idealized version of concept, which in our case, is modelled with a learned vector.
2

---

attention-bass algorithms such as Recurrent Independent Mechanisms (RIMs) and its derivatives (23 24) the use of discrete modularization of features. But, unlike all these other contemporary
algorithm, ConSpec aims to directly focus on identifying critical steps, a less burdensome task than the typical modelling of value, reward, or actions taken for all encountered states.
ConSpec was inspired, in part, by the contrastive learning literature in computer vision (26 13 61 27). Specifically, it was inspired by the principle that transforming a task into a classification problem and using well-chosen positive and negative examples can be a very powerful means of learning invariant and generalizable representations. A similar insight on the power of classification led to a recent proposal for learning affordances (32), which are arguably components of the environment that are required for critical steps. As such, ConSpec shares deep connection with this work. Along similar lines, (56) 2 [54 [9] 21;;
T8 36 D have begun to explore contrastive systems and binary classifiers to do imitation learning and RI.. With these works, ConSpec shares the principle that transforming
the RI. problem into classification task is beneficial. But these works typically use the contrastive approach to learn whole policies or value functions. distinction to this, ConSpec uses contrastive learning for the purpose of learning prototype representations for recognizing critical steps in order to shape the reward function, thereby enabling rapid learning and generalization.
2.1 Neuroscience inspiration
Our design of ConSpec was ultimately inspired at high level by theories from neuroscience. For example, it shares with episodic control and related algorithms (37 8 45 40 43) the principle that memory bank can be exploited for fast learning. But unlike episodic control, ConSpec exploits memory not for estimating a surrogate value function, but rather, for learning prototypes that hone in on small number of discrete critical steps. In-line with this, ConSpec takes inspiration from the neuroscience literature on the hippocampus, which suggests that the brain, too, endeavours to encode and organize episodic experience around discrete critical steps (57) 63 64 20). Interestingly, recent studies suggest that honing in on discrete critical states for success may in fact be key to how the brain engages in reinforcement learning altogether (2)). Contemporary RI. is predicated on the Bellman equation and the idea of reward prediction. But recent evidence suggest that dopamine circuits in the brain, like ConSpec, focuses retrospectively causes of rewards (29). i.e. critical steps.
3 Description and intuition for ConSpec
Why the use of contrastive learning in ConSpec? Humans are very good, and often very fast, at recognizing when their success is contingent upon multiple past steps. Even children can identify contingencies in a task after only handful of experiences (55) 22) Intuitively, humans achieve this by engaging in retrospection to examine their past and determine which steps along the way were necessary for success. Such retrospection often seems to involve comparison: we seem to be good at identifying the differences between situations that led to success and those that led to failure, and then we hone in on these differences. We reasoned that it would be possible for RL agents to learn to identify critical steps-and and moreover, to do so robustly-ifs a system were equipped with a similar, contrastive, capability to take advantage of memory sort through the events that distinguished successful versus failed episodes. This paper's 's principal contribution is the introduction of an add-on architecture and loss for retrospection of this sort (Fig. D.
At its core, ConSpec is a mechanism by which a contrastive loss enables rapid learning of a set of prototypes for recognizing critical steps. Below describe how ConSpec works.
3.1 A set of invariant prototypes for critical steps
To recognize when a new observation corresponds to a critical step, we store a set of H prototypes in a learned representation space. Each prototype, h; is a learned vector that is compared to non-linear projection of the latent representation of currently encoded features in the environment, go go(2)) (with projection parameters 8). Thus, if we have current observation Oz, encoded as latent state 2, = fw (O.) (with encoder parameters W), we compare gu(22) to each of the prototypes, h, in order assess whether the current observation corresponds
critical step or not.
At first glance, one may think that a large number of prototypes is needed, since there are a massive number of states corresponding to critical steps in non-toy tasks. But, by using cosine similarity
3

---

between prototype h2 and encoder output 2g, cos (h,, 90(2)), ConSpec can recognize a large, poten- tially infinite set of states as being critical steps, even with small number of prototypes, because the mapping from O₂ to cos (hx, 30 (=+) is many-to-one mapping, akin to neural net classifier.
In what follows, used H 20 for all tasks, continuous control, which have lot of state-level variation. Of course, in other settings more prototypes may need be used. In general, one can specify more prototypes than necessary, which does not hurt learning (Fig. A.3 for an extreme example).
the ones in 3D environments and with
3.2 A memory system for storing successes and failures
ConSpec maintains three memory buffers: one for the current mini-batch being trained on (B; with capacity of MB), one for successful episodes (S; with a capacity of Ms), and one for failure episodes (F; with capacity of My). Each episode is of length T time steps. As we show, these memory buffers need not be very large for ConSpec to work; we find that 16 slots in memory is enough, even for 3D environments. Each buffer stores raw observations, Os, encountered by the agent in the environment. When a new mini-batch of observations loaded into B for training, the episodes are categorized into successes and failures using the given criteria for the task. We then fill S and F with observations from those categorized episodes in intrstttitttt (FIFO) manner, where episodes stored from previous batches overwritten as new trajectories come in.
3.3 A contrastive loss to learn prototypes
To learn our invariant prototypes of critical steps (and the encoding of the observation) we employ contrastive loss that differentiates successes from failures. The contrastive loss uses the observations stored in S and F to shape the prototypes such that they capture invariant features shared by a spectrum of observations in successful episodes. The functional form of this loss is:
H
$1
LConSpec
Σ
1
Σ1
11
max SEKel + AE(I...T)
max strl+
Ms
Mr
TE(1...T)
1=1
RES
REF
Σ
D(f hisiH)
H
KES
where Sikr cos (h,, go (=ke)) the cosine similarity of prototype h2 and the latent representation for the observation from timestep found in episode k, Sik is the vector of sik elements concatenated along t and softmaxed over t, and D(-) is function that measures the diversity in the matches the prototypes (see Appendix section A.5 for functions (--) we used in this work). Note the purpose of the softmax was to sparsify the sequence of cosine scores for each trajectory over the stee-steps! t. Our goal was to compare pairs of prototypes and force their peak cosine scores to be different in time. The softmax was one way to do this, but any equivalent method would equally well.
To elaborate on this loss, for each prototype h2 and each episode's projection at time step 1, go (zke). the cosine similarity normalized dot product) is calculated between the prototype's vector and the latent projection vectors to yield a time-series of cosine similarity scores. The maximum similarity score across time within each episode is then calculated. The maximum similarity score for each successful episode is pushed up (Ist term of the loss) and for each failed episode, it is pushed down (2nd term). The 3rd term of the loss encourages the different prototypes to be distinct from one another. The ConSpec loss is added to any RL losses during training (Algorithm I pseudocode). Below we show that this loss leads to prototypes that hone on distinct critical steps (Fig. 44).
There is a potential nonstationarity in the policy that could cause instability in the prototypes discovered, the policy improves and previous successes become considered failures. To prevent this, we froze the set of memories that defined each prototype upon reaching the criterion defined in Appendix A.5 In practice, surprisingly, ConSpec does not suffer from learning instability in the multikey-to- -to-oor task even if prototype freezing is not done (see Fig. A.12) but it is imaginable that other tasks may still suffer so the freezing mechanism was kept.
3.4 Generating intrinsic rewards based on retrospection
Thus far, the described model learns a set of prototypes for critical steps along with an encoder. This system must still be connected to RL agent that can leam policy (Tz. parameterized by

---

Figure 1: ConSpec is mod- ule that learns hone in on critical state. ConSpec trains its prototypes by comparing successful v. failed episodes via a contrastive loss, and learning incentivizes pushing cosine similarity scores of suc- cesses to 1, and failures to O. It then uses the match to these prototypes to output intrinsic rewards to the RL agent.
ConSpec
Aie
cari
oel
... Bosc
Go
Av
Encader
Episode buffers
Backbone system
Palkyy
Lya
8
Environment
)). This could proceed in a number of different ways. Here, we chose the reward shaping approach
522 47 42), wherein we add additional intrinsic rewards tRz to the rewards from the environment, Tits though other strategies are possible. One simple approach is to define the intrinsic reward episode k at time-step FRt. as the maximum prototype score, 8ikt, across moving window of H time-steps (H in our experiments, with threshold of 0.6 applied to ignore small score fluctuations), and scaled to be in the correct range of the task's s expected nonzero reward per step Rtask:
H
kkt A Reask Σ
Sike 1as (Bik,t (B.... 81..11)
However, concern with any reward shaping is that it can alter the underlying optimal policy (47 42) Therefore, we experimented with another scheme for defining intrinsic rewards:
H (y
FRr A Prask
Sak,t 8ik,t-1)
where À is proportionality constant and 7 is the discount in RL backbone. This formula satisfies the necessary and sufficient criterion from (42) for policy invariance and provably does not alter the optimal policy of the RL agent. In practice, we find that both forms of intrinsic reward work well (Fig. and A.II so unless otherwise noted, results in this paper used the simpler implementation, Eqn.
3.5 Inductive biases in ConSpec that enable rapid eredit assignment
What is the intuition for why ConSpec's inductive biases make it well suited to learn when success is contingent upon multiple critical steps? Learning a full model of the world and its transitions and values is very difficult, especially for novice agent. Most other RL approaches can be characterized, at a high level, as learning to predict rewards or probability of success given sequences of observations in memory. In contrast, ConSpec solves the reverse problem of retrospectively predicting the critical steps given success. This is much easier problem, because in many situations, critical steps are independent of one another when conditioned on success. As such, when we consider the joint distribution between success and all of the critical steps, we can learn it retrospectively by considering the conditional probability given success for each potential critical step individually. In contrast, if we learn the joint distribution prospectively, then we must learn the conditional probability of success given all the potential critical steps tandem, which a much more complicated inference problem. To see concrete example, conditioned on the successful acceptance of paper, one be sure that all the critical steps had each individually been achieved (an idea conceived, experiments run, paper written and reviewer concerns addressed). Given that all of the prototypes are learned in parallel, this means that the search for each of these critical steps do not depend on one another when taking this retrospective approach. This is large gain over any algorithm attempting to learn the function of how combinations of steps could predict reward.
4 Empirical results
We tested ConSpec on a variety of RL. tasks: grid worlds, continuous control, video games, and 3-D environments. Training was done on an RTX 8000 GPU cluster and hyperparameter choices are de- tailed in the A.5. We show how ConSpec rapidly hones on critical steps, helping alleviate two difficult RL problems: long term credit assignment, and out-of-distribution generalization.
5

---

Algorithm/ 1 Jointly training ConSpec with an RL agent
Given: Current parameters (W, O, ④. h)...)) ) and Inputs: memory buffers B, S, and F, number of episodes in a mini-batch, B, and number of epochs to train for, E
1: for Epoch 1...E do
2: Collect new minibatch of B trajectories using Th and store them {(OI:T (1: 01:7, (1rr)) in B
3: Update the S and F memory banks based on B
4: Calculate latent representations: 2kt <- f(Ort)
5:
Calculate scores: Sikt <- cos (ha, 9o (=k+))
6:
Calculate intrinsic rewards, F'ke and total rewards, Tkt
Total Fkt + F'ke Calculate RL loss, LRI., (per RL backbone) and ConsSpec loss per equation
7:
8:
Update parameters via loss Lioal = LRI. + BLoonspes
9: end for
Figure 2: ConSpec learns invari- ant representations. (a) 3D Orange- Tree task design. (b) ConSpec effi- ciently learns this task approaching the maximum reward (=10), while PPO, SynthRs. and CURL base- lines found it difficult. (c) This pro- totype learns to robustly detect the retrieval of flat objects, invariant variations in object type and colour.
Orange-Tree Task
Max reward
Canápec
Revorn
140
CUR.
Starch nat object
Pick nat onject
Waiting room
Cimb foat cbject
Pick orange
Minibatches
Example preferred state in & differant apisodas by prototypa hi
4.1 ConSpec learns prototypes for critical steps in 3D environments
We began with series of experiments in 3D virtual environment. This was verify that Con- Spec can identify critical steps in complex sensory environments, wherein critical steps do not correspond to specific observable states (as bottleneck states do). To do this, used SilicoLabs [ Unity-based (31) game engine to create a 3-stage task that we call the "OrangeTree" task. stage T of this task, the agent is placed in room with two objects, one of them round spiky (e.g. balls jacks), and one of them flat (e.g. boxes or tables) (Fig. 2b) and the agent can pick one them up. In stage 2, the agent is put in waiting room where nothing happens. In stage 3, the agent is placed in a new room with an orange tree. To obtain a reward, the agent must pick an orange, but to do this, must stand on flat object (which it can do only if it had picked up flat object in stage I). So in stage 1, the agent must learn the critical step of picking up a flat object. But, because there are a variety of flat objects, there is no single bottleneck state to learn, but rather, a diverse set of states sharing an invariant property (i.e. object flatness).
ConSpec was compared with series of baselines. We first tested Proximal Policy Optimization (PPO) 53 33 on this task, which failed (Fig. (D). One possibility for this is the challenge of long-term credit assignment, which arises from the use of the waiting room in stage 2 of the task. As such, we next tested an agent trained with Synthetic Returns (SynthRs) (48) atop PPO, a recent reward shaping algorithm designed solve long-term credit assignment problems. SynthRs works by learning to predict rewards based on past memories and delivering an intrinsic reward whenever the current state is predictive of reward. But SynthRs did not help solve the problem (Fig. P showing that the capacity to do long-term credit assignment is, by itself, not enough when success does not depend on single bottleneck state but rather a diverse set of states that correspond to a critical step.
Thus, we next considered more sophisticated representation learning for RL: Contrastive Unsu- pervised Reinforcement Learning (CURL) (56), which outperforms prior pixel-based methods on complex tasks. However, adding CURL to PPO, too, did not solve the task (Fig. (p)). likely because CURL is not learning invariances for task success, but rather, invariances for image perturbations.
Finally, we examined the performance of ConSpec, An agent trained by PPO with additional intrinsic rewards from a ConSpec module (per Algorithm[[] can indeed solve the OrangeTrree task (Fig. 2b). To further dissect the reasons for ConSpec's 's abilities, we studied the invariances learned. To do this, we collected successful episodes and then identified the range of observations (images) that maximally matched the prototypes during the episodes. Notably, one of the prototypes consistently
https//wwww.ilicolasscca

---

TrangeTree black objects
b
Preferred state in apisedes prototype
Figure 3: ConSpec's invariant represen- tations help OoD generalization. (a, c) New Orange'Tree tasks during Testing (a) with black boxes/tables, or (c) in differ- ently coloured room. Previously trained prototypes hone in on interpretable states even in these new contingencies (b,d) and are able to generalize (e) in zero- shot with never-before seen black ob- jects, and (f) in few-shot in a new en- vironment, approaching the max reward.
Triiring
Teating
OrangeTroc different contect
Preferred state in apisodes by prototype
Training
Testing
Man renard
Mas rwward
Rewerd
Senthes IR0 PPD
10 15
Contiee
Grad
had observations of picking up a flat object mapped to it. In other words, this prototype had discovered an invariant feature of critical step that was common to all successes but not failures, namely, the step of having picked up flat object (Fig. 2 Moreover, this prototype learned to prefer flat objects regardless of the type of object (i.e. box or table) and regardless of its colour. So, ConSpec's 's contrastive learning permitted the prototype to learn to recognize a critical step, namely picking up flat objects, and maintain invariance to irrelevant features of those objects such as shape and colour.
4.2 ConSpec generalizes to new contingencies in 3D environments
We next asked whether ConSpec's capacity for finding invariances among critical steps could aid in zero-shot generalization when the sensory features of the environment were altered. To test this, we made a variation of the Orange Tree task. During training. the agent saw flat objects that were magenta, blue, red, green, yellow, or peach (Fig. BR.. During testing. a black table or box was presented, even though the agent had never seen black objects before. The prototype that recognized picking up flat objects immediately generalized to black flat objects (Fig. 3))) Thanks to this, the agent with ConSpec solved the task with the black objects in zero-shot (Fig. 3e. This shows that ConSpec was able to learn prototypes that discovered an invariant feature among successes (flatness), and was not confused by irrelevant sensory features such as colour, permitting it to generalize to contingencies involving colours never seen before.
To further test ConSpec's capacity for generalization, we made another variation of the Orange task with the background environment changed, in a manner typical of -fitt--diitiibutiin (OoD) generalization tests (4). In particular, testing took place in a different room than training, such that the new room had pink walls and green floor (unlike the gray walls and blue floor of the training room) (Fig. Be). Again, the prototypes generalized in a zero-shot manner, with flat objects being mapped to the appropriate prototype in the new environment despite the different background (Fig. Bt... The policy network did not immediately generalize, but after brief period of acclimatization the new environment (where the underlying PPO agent was trained for no more than 20 gradient descent steps while ConSpec was kept constant), the ConSpec agent was able to rapidly solve the new but related task environment (Fig. 3)
4.3 ConSpec improves credit assignment in gridworld tasks with multiple critical steps
To investigate more deeply ConSpec's ability to engage in long-term credit assignment by identifying critical steps, we used simple
tasks in 2D grid-worlds, These tasks were meant to capture in the simplest possible instantiation the idea that in real life, rewards are sparse and often dependent on multiple critical steps separated in time (akin to the super-acceptance task structure tasks we used here, the agent must find one or more keys to open one or more doors (Fig. 4.). If the agent does not pick up all of the keys and does not pass through all of the doors, success is impossible. To make the tasks challenging for Bellman back-up, the agent is forced to wait a wait room for many time steps in-between the retrieval of each key. An episode is successful if and only if the agent exits the final door.
described in the Introduction D. In the malti-kee-tt-ddool
We found that when ConSpec was trained these tasks, each prototype learned to represent a distinct critical key-retrieval step (e.g. one prototype learned to hone in on the retrieval of the first key, another prototype on the second key, etc.). We demonstrated this by examining which state in different episodes maximally matched each prototype. As can be seen, picking up the first key was
7

---

Key
Key 2 Ooor
Key
Kay
Wetting
Walting
Watting
Walting
Prototype h
Prototype
Prototype fta
Prototype hn
G
C
Highest matching state in example episodes for each prototype hy. by. byy
ConSipec
ConSipcc minus
d
TVT
expforation time
Roward
keys kees
keys
kays
keys
4000
8000
4000
8000
4000
SynthRs
RUDDER
10
10
keys
keys
bays
keys
Jeys
3 koys
loys
0000
4000
4000
8000
Minibatches
Minibatches
Minibatches
Figure 4: ConSpec rapidly learns tasks with multiple contingencies akin to the scenario from the Introduction D (a) Protocol for key-tt--otoor task with 4 keys. (b) Plotted: states that maximally matched each prototype. Here, prototypes learn to hone in on states that depict retrieval of the key (where the agent is out the door and the coloured keys are gone i.e. retrieved). (c left) ConSpec rapidly learns the multi-key-to-door tasks, whereas (d-g) TVT, SynthRs, RND, and RUDDER performances collapse. (c right) When exploration time is subtracted, ConSpec's 's training time x Occonsaant) approximately, even as number of keys increased, affirming the complexity predictions from 3.5
matched to one prototype, picking up the second was matched to another, and so on (Fig. 4p). Thus, ConSpec identifies multiple critical steps in a highly interpretable manner in these tasks.
We then tested ConSpec atop PPO again. We compared performance to four other baselines. The first three, SynthRs, Temporal Value Transport (TVT) system (28), and RUDDER (3) are contemporary reward-shaping based solutions to long-term credit assignment in RL. The fourth baseline was random network distillation (RND) (10) a powerful exploration algorithm for sparse reward tasks, in order to study exploration itself was sufficient to overcome credit assignment in situations with multiple critical steps.
We found that PPO with ConSpec learns the multi-key-to-due tasks very rapidly, across a range of number of keys (Fig. 4c. left). Importantly, even as keys got added, ConSpec's 's training time remained constant when the exploration required for new keys was accounted for (Fig. 4.. right), consistent with the intuitions given above (section 3.5)). On the other hand, RUDDER, SynthRs, TVT, and RND solved the task for a single key thanks their ability to do long-term credit assignment (Fig. A.2] A.4 A.5) but their performance collapsed more keys were added (Fig. 41-g), highlighting the difficutty long-term credit assignment with multiple critical steps.
We also note that RND, despite being strong enough to handle sparse reward tasks like Montezuma's Revenge (10), also fails here as more keys are added, illustrating the point that long term credit assignment is its own distinct issue beyond exploration which ConSpec addresses and RND does not.
What happens when the number of prototypes is not sufficient to cover all the critical steps? Even having fewer than necessary prototypes (3 prototypes in the 4-key task, shown in Fig. A.8) can surprisingly still be enough to solve the task (i.e. catching any critical step at all, still helps the agent).
also tested ConSpec on another set of grid-world tasks with multiple contingencies, but with different and harder to learn task structure. Again, ConSpec but not other baselines could solve these tasks (Fig. 16),), illustrating that ConSpee successfully handles a wide variety of real-life inspired scenarios with multiple contingencies.
8

---

Finally, we note that ConSpec successfully learns, not only atop a PPO backbone, but also atop an Reconstructive Memory Agent (RMA) backbone (Fig. A.10 and also Fig. A.2p), illustrating ConSpec's flexibility in being added to any RI. agent to improve long-term credit assignment.
4.4 ConSpec helps credit assignment in delayed Atari tasks and Montezuma's 's Revenge
Next studied how ConSpec could help credit assignment in other RL settings, such as the Atari games from OpenAI Gym (), including in some unanticipated settings where it would not have been expected to help. We first note that existing RL algorithms, such as PPO, can solve many of these games, as does PPO with ConSpec (Fig. A.17 But, PPO alone cannot solve delayed reward versions of the Atari games. These require specialized long-term credit assignment algorithms solve them And adding ConSpec to PPO also solves delayed Atari (Fig. 5
3 48
Delayed Pong
b
Delayed Bowling
Delayed AirRaid
Figure 5: ConSpec improves per- formance on delayed (a) Pong, (b) Bowling, and (c) AirRaid.
Toward
20
2000
4050
2000 Masibuichca
4020
2000 Va Hi-Bbatshca
4000
ConSome
Within the set of unmodified Atari games, PPO alone also does not solve Montezuma's Revenge, and it does not usually exceed 0 reward, so we investigated this game further. Montezuma's 's Revenge is usually considered to be an exploration challenge, and sophisticated exploration algorithms have been developed that solve it like RND(1O) and Go-Explore (16). But, interestingly, we that adding ConSpec to PPO improves performance on Montezuma's 's Revenge even without any sophisticated exploration techniques added (Fig. 6b). To provide intuition as why ConSpec is able help even in the absence of special exploration, we studied how many successful trajectories takes for ConSpec to learn how obtain at least the first reward in the game. Crucially, ConSpec atop PPO with only a simple e-greedy exploration system is able to learn with as few as 2 successful episodes (Fig. (p). We speculate that this is because ConSpee, like other contrastive techniques (33), benefits mostly from a large number of negative samples (i.e. failures) which are easy to obtain. Other algorithms like Recurrent Independent Mechanisms (RIMs) have demonstrated large improvements on other Atari games (23), and make use of discrete modularization of features as ConSpec does by its prototypes. But RIMS discretizes by very different mechanism, and does not make progress on Montezuma's 's Revenge (Fig. (C).
Interestingly, Decision Transformers (DT) are designed learn from successful demonstrations as ConSpec does (Fig. Gd) but cannot solve Montezuma's 's Revenge if those successes were "spurious successes" (an important learning signal for novice agents) during random policies, affirming that DTs, unlike ConSpec, need to be behaviourally cloned from curated experts. This was still true when the DT's were trained with every success trajectory uniformly labelled +1 and every failure O. Even with this equally weighting of data, DTTs still got O reward. (Fig. (6)) By contrast, ConSpec learns well by inherently treating the data unequally, honing in on only few critical steps.
Putting this altogether, the secret to ConSpec's 's rapid abilities to learn from spurious successes arises because the prototypes ignore irrelevant states and learn to hone in the few critical ones, allowing ConSpecto get reward on challenging tasks like Montezuma's 's Revenge, doing so even without dedicated exploration algorithm. We speculate that coupled with an exploration algorithm ConSpec could be even more powerful.
4.5 ConSpec helps credit assignment in continuous control tasks
Finally, we study another scenario in which ConSpec uncxpectedly helps. Although ConSpec was designed to exploit situations with discrete critical steps, we studied whether ConSpec might help even in tasks that do not explicitly involve discrete steps, such as continuous control tasks. These tasks are interesting because both the robotics and neurophysiology literature has long known that "atomic moves" can be extracted from the structure of continuous control policies (15) 60). Moreover, this is evident in real life. For example, learning to play the piano is continuous control task, but we can still identify specific critical steps that are required (such as playing arpeggios smoothly). However, contemporary RI. algorithms for continual control tasks are not designed to hone in on sparse critical steps, and instead, they focus on learning the Markovian state-to-state policy.

---

Conspec Montezuma's revenge
Reward
400
ConSrec
400
200
200
1000
2000
2000
800
1600
Minibatches
Minibatches
Reward n
RIMs
d
DTransformers reward-to-go
DTransformerr,
400
400
400
uniformly labelled
200
200
DTaanfomer successes
200
1000
2000
3000
1000
2000
3000
1000
2000
3000
Minibatches
# Minibatches
# Minibatches
Figure 6: (a) PPO+ConSpec progresses on Montezuma's 's Revenge despite the lack of a dedicated exploration algorithm. (b) PPO+ConSper with 2, 4, and success demonstrations, highlighting how ConSpec uses few successes very efficiently. (c) RIMs, and (d) DT's get 0 reward, even when DT's are given up 36 spurious success demonstrations. (e) DT's trained with success trajectories uniformly labelled +1 and failures O. With this uniform weighting of the data, DT's still get 0 reward.
We focused on delayed versions of three continuous control tasks: Half Cheetah, Walker, and Hopper, where the delay was implemented according to past literature (49) introduce long-term credit challenge (see A.4). ConSpec atop PPO is able to significantly improve performance in all tasks (Fig. showing the utility of learning identify critical steps even in continuous control tasks.
Delayed Hopper
Delayed Walker
Delayed Half Cheetah
Figure 7: ConSpec improves per- formance in delayed versions of continuous control tasks (a) Hop- per, (b) Walker, (c) Half Chectah.
Reward
s00 800
500
200
100
-100
40
0.50
500
1000
1000
630
1000
minibatches
FPO
minibotches
minibatches
+PPO
Discussion and limitations
Here, we introduce ConSpec, contrastive learning system that can be added to RL agents to help hone in on critical steps from sparse rewards. ConSpec works by learning prototypes of critical steps with novel contrastive loss that differentiates successful episodes from failed episodes. This helps both rapid long-term credit assignment and generalization, as we show. Altogether, ConSpec is a powerful general add-on to any RL agent, allowing them to solve tasks they are otherwise incapable.
Despite its benefits for rapid credit assignment and generalization beyond Bellman- saatmannbackuu-baaee based approaches in RL, ConSpec has limitations, which we leave for future research. For one, the number of prototypes is an unspecified hyperparameter.
our experiments, we did not use more than 20, even for rich 3D environments, and in general one can specify more prototypes than necessary, which does not hurt learning (Fig. A..3) for an extreme case). However, an interesting question is whether change needed.
it would be possible to design system where the number of prototypes For another limitation, ConSpec requires a definition of success and failure. We used relatively simple definitions for tasks with sparse rewards well for tasks with dense rewards as detailed in section A.3] and they worked well. But, future work can determine more principled definitions. be learned conditioned, and as such, are related to the
Alternatively, definitions of success topic of goal-conditioned RL, which is left unexplored here. Nonetheless, we believe that this is promising future direction for ConSpec, and that it may potentially be of great help in the yet unsolved problem of sub-goal and options discovery.
6
Acknowledgements
We heartily thank Doina Precup, Tim Lillicrap, Anirudh Goyal, Emmanuel Bengio, Kolya Malkin, Moksh Jain, Shahab Bakhtiari, Nishanth Anand, Olivier Codol, David Yu-Tung Hui, well as members of the Richards and Bengio labs for generous feedback. We acknowledge funding from the Banting Postdoctoral Fellowship, NSERC, CIFAR (Canada AI Chair and Learning in Machines & Brains Fellowship), Samsung, IBM, and Microsoft, This research was enabled in part by support provided by (Calcul Québec) and the Digital Research Alliance of Canada.
10

---

References
[1] Agarwal, R., Machado, M. C., Castro, S., and Bellemare, M. Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning. arXiv e-prints, art. arXiv:2101.05265, January 2021.
[2] Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-.,, and Devon Hjelm, R. Unsupervised State Representation Learning in Atari. arXiv e-prints, art. urXiv:1906.088. June 2019.
[3] Arjona-Medina, J. A., Gillhofer, M., Widrich, M., Unterthiner, T., Brandstetter, J., and Hochreiter, S. RUDDER: Return Decomposition for Delayed Rewards. arXiv e-prints, art. arXiv:1806.07857, June 2018.
[4] Arjovsky, M. Out of Distribution Generalization in Machine Learning. arXiv e-prints, art. arXiv:2103.02667, March 2021.
[5] Bacon, P.-L., Harb, J., and Precup, D. The Option-Critic Architecture. arXiv e-prints, art. arXiv:1609.05140 September 2016.
[6] Barhate, N. Minimal implementation of decision transformer. https://github.com com/
transforrrrttrrrfrrrrrrro 2022.
[7] Bellman, R. Dynamic programming. Technical report,Princente, University Press, 1957.
[8] Blundell, C., Uria, B., Pritzel, A., Li, Y., Ruderman, A., Leibo, J. Z., Rae, J., Wierstra, D., and
Hassabis, D. Model-Free Episodic Control. arXiv e-prints, art. arXiv:1606.044: 04460, June 2016.
[9] Brockman, G. Cheung. V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba. W. Openai gym, 2016.
[10] Burda, Y., Edwards, H., Storkey, A., and Klimov, O. Exploration by Random Network Distillation. arXiv e-prints, art. arXiv:1810. 12894, October 2018. doi: 10.48350/arXiv.1811 12894.
[11] Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. https://gt giigttth..... com/kaa//ecission wrwnssoonnttassommere..... 2020.
[12] Chen, L.., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, Decision Transformer: Reinforcement Learning via Sequence Modeling. arXiv e-prints, art. arXiv:2106.01345. June 2021.
[13] Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A Simple Framework for Contrastive Learning of Visual Representations. arXiv e-prints, art. arXiv:2002.05709, February 2020.
[14] Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. Leveraging Procedural Generation to Benchmark Reinforcement Learning. arXiv e-prints, art. arXiv: December 2019.
[15] d' A'eell,, A. and Bizzi, E. Shared and specific muscle synergies in natural motor behaviors. Proceedings of the National Academy of Sciences, 102 102(8) 2005. doi: 10.1073/pnas. 0500199102. httpsttp: //www//wwww/////// pnas org /oi abs 10. 1ww100077aaaa 0500199102
[16] Ecoffet, A., Huizinga, J., Lehman, J., Stanley, K. O., and Clune, J. First return, then explore. Nature, 590(7847):580-588) February 2021. doi: 10.10038/441586-2--011333
[17] Espeholt, L.,, Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning. I., Legg, S., and Kavukcuoglu, K. IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures. arXiv e-prints, art, arXiv: arXiv:1802.0156) 180.001561, February 2018.
[18] Eysenbach, B., Zhang, T., Salakhutdinov, R., and Levine, S. Contrastive Learning as Goal- Conditioned Reinforcement Learning. arXiv e-prints, art. arXiv:2206.07568, June 2022.
[19] Finn, C., Christiano, P., Abbeel, P., and Levine, S. A Connection between Generative Adversarial Networks, Inverse Reinforcement Learning, and Energy-Band Based Models, arXiv e-prints, art. (XXi:16111.03855 November 2016.

---

[20] Franklin, N. T., Norman, K. A., Ranganath, C., Zacks, J. M., and Gershman, S. J. Structured event memory: a neuro-symbolic model of event cognition. bioRxiv, 2019. doi: 10.1101/541607. http:htts://ww. http://wwwwborxiivco/o/.// www.biiooiii........o wtwlooiiiorgg/oooooooooooo http://ttteeaaa////////////00000
[21] Ghosh, D., Gupta, A., Reddy, A., Fu, J., Devin, C., Bysenbach, B., and Levine, S. Learning to Reach Goals via Iterated Supervised Learning. arXiv e-prints, art. 1rX.v:1912.06088, December 2019.
[22] Gopnik, A. and Sobel, D. M. Detecting blickets: How young children use information about novel causal powers in categorization and induction. Child Development, 71(5):1205-13 1222, 2000. ISSN 00093920, 14678624. Utt.http:/ htttt///// http:://wwww...orrooroo/////
[23] Goyal, A., Lamb, A., Hoffmann, J., Sodhani, S., Levine, S., Bengio. Y., and Schölkopf, B. Recurrent Independent Mechanisms. arXiv e-prints, art. arXiv:1909, 10893, September 2019.
[24] Goyal, A., Lamb, A., Gampa, P., Beaudoin, P., Levine, S., Blundell, C., Bengio, Y., and Mozer, M. Object Files and Schemata: Factorizing Declarative and Procedural Knowledge in Dynamical Systems. arXiv e-prints, art. utXiv:2006.1622, June 2020.
[25] Goyal, A., Lamb, A., Hoffmann, J., Sodhani, S., Levine, S., Bengio, Y., and Schölkopf, B. Recurrent independent mechanisms.
com/ wwom/anniii www.annnnn/////aI.. 2020.
[26] Hadsell, R., Chopra, S., and LeCun, Y. Dimensionality reduction by learning invariant mapping. In 2006 IEEE Computer Sociery Conference on Computer Vision and Pattern Recognition (CVPR (06), volume 2. pp. 1735-1742, 2006. doi: 10.1 10.1109/CVP..2006... 100.
[27] He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Momentum Contrast for Unsupervised Visual Representation Learning. arXiv e-prints, art. 19X1::1911.05722 November 2019.
[28] Hung, C., Lillierap, T. P., Abramson, J., Wu, Y., Mirza, M., Carnevale, F,, Ahuja, A., and Wayne, G. Optimizing agent behavior over long time scales by transporting value, Nat Commun, 10, 2019. doi: https://doi.org/ 00. 1038/44
[29] Jeong, H., Taylor, A., Floeder, J. R., Lohmann, M.., Mihalas, S., Wu, B., Zhou, M., Burke, D. A., and Namboodiri, V. M. K. Mesolimbic dopamine release conveys causal associations. Science, x7866626), 2022. doi: /ttp////eecce.agg//4444 URL htttps://www.sciii org/doi/abae http:/11.11////cienneee abq6700
[30] joo, C. random- innatom--nttorr-ddiiiillaooooooo
madom--etwwat--dssilliioooyyoo 2018.
https://github..co////////
2018.
[31] Juliani, A., Berges, V.-P., Teng, E., Cohen, A., Harper, J., Elion, C., Goy, C., Gao, Y., Henry, H., Mattar, M., and Lange, D. Unity: A General Platform for Intelligent Agents. arXiv e-prints, art. arXiv:1809 809.026270 September 2018.
[32] Khetarpal, K., Ahmed, Z., Comanici, G., Abel, D., and Precup, D. What can I do here? A Theory of Affordances in Reinforcement Learning. arXiv e-prints, art. arXiv:2006.15083 June 2020. doi: 10.48550/arxiv.2006..55 15085.
[33] Kostrikov, I.
of reinforcement learning algorithms. https://github com ctwp://ooorrrooooyrrrra------eeo gaiil 2018.
[34] Laskin, M. Curl: Contrastive unwpbiirveduuuurriiiiiseeeel representation learning for sample-efficient rein- forcement learning. https://githbb. http://gtttttacooom///saaaaaaaaa 2022.
[35] Lengyel, M. and Dayan, P. Hippocampal Inf. Process. Syst., volume 20, 11 2007.
conpeccaappcoootrtttttttoioo contributions to control: The third way. In Adv. Neural
[36] Li, B., François-Lavet, V., Doan, T., and Pineau, J. Domain Adversarial Reinforcement Learning. arXiv e-prints, art. arxiv:2102.00009, February 2021.
[37] Lin, I. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine Learning, 8:293-321, 2004.
12

---

[38] McGovern, A. and Barto, A. G. Automatic discovery of subgoals in reinforcement learning using diverse density. In Proceedings of the Eighteenth International Conference on Machine Learning. ICML '01. pp. 301-3611,,, San Francisco, CA, USA, 2001. Morgan Kaufmann Publishers Inc. ISBN 1558607781.
[39] Mnih, V. Kavukeuoglu, K., Silver, D., Graves, A., Antonoglou, I., Wierstra, D., and Riedmiller, M. Playing Atari with Deep Reinforcement Learning. arXiv e-prints, art. December 2013.
[40] Moore, A. and Atkeson, Memory-based reinforcement learning: Efficieet computation with prioritized sweeping. In Hanson, S., Cowan, J., and Giles, C. (eds.), Advances in Neural Informa- fion Processsig Systems, volume Morgan-Kaufmanne 1992. UKL.httts::/proceddiii pdf
[41] Nachum, O., Gu, S., Lee, H., and Levine, S. Data-Efficient Hierarchical Reinforcement Learning. arXiv e-prints, art. urXiv:180500826, May 2018.
[42] Ng, Y., Harada, D., and Russell, Policy invariance under reward transformations: Theory and application to reward shaping. In In Proceedings of the Sixteenth International Conference on Machine Learning, pp. 278-287. Morgan Kaufmann, 1999.
[43] Oh, J., Guo, Y., Singh, S., and Lee, H. Self-Imitation Learning. arXiv e-prints, art. 1.Xi::8006006666 June 2018
[44] Patil, V. P., Hofmarcher, M., Dinu, M.-C., Dorfer, M., Blies, P. M., Brandstetter, J., Arjona- Medina, A., and Hochreiter, S. Align-RUDDER: Learning From Few Demonstrations by Reward Redistribution. arXiv e-prints, art. 1Xiv:2009.14108, September 2020. doi: 10-48550/arXi..2009.444
[45] Pritzel, A., Uria, B., Srinivasan, S., Puigdomènech, A., Vinyals, O., Hassabis, D., Wierstra, D., and Blundell, C. Neural Episodic Control. arXiv e-prints, art. arXiv: March 2017.
[46] Puigdomènech Badia, A., Piot, B., Kapturowski, S., Sprechmann, P., Vitvitskyi, A., Guo, D., and Blundell, C. Agent57: Outperforming the Atari Human Benchmark. arXiv e-prints, art. arXiv:2003.155 13350, March 2020.
[47] Randløv, J. and Alstrom, Learning to drive bicycle using reinforcement learning and shaping. In Proceedings of the Fifteenth International Conference Machine Learning, ICML '98, pp. 463-471, San Francisco, CA, USA, 1998. Morgan Kaufmann Publishers Inc. ISBN 1558605568.
[48] Raposo, D., Ritter, S., Santoro, A., Wayne, G., Weber, T., Botvinick, M., van Hasselt, H., and Song, F. Synthetic Returns for Long-Term Credit Assignment. arXiv e-prints, art. arXiv:2102. 12425, February 2021.
[49] Rengarajan, D., Vaidya, G., Sarvesh, A., Kalathil, D., and Shakkottai, S. Reinforcement Learning with Sparse Rewards using Guidance from Offline Demonstration. arXiv e-prints, art. arXiv:2202.046288 February 2022. doi: 00448550/rXx..2222.44000
[50] Sasha Vezhnevets, A., Osindero, S., Schaul, T., Heess, N., Jaderberg, M., Silver, D., and Kavukcuoglu, K. FeUdal Networks for Hierarchical Reinforcement Learning. arXiv e-prints, art. urXiv:1703. 01161, March 2017.
[51] Schaul, T., Quan, J., Antonoglou, I., and Silver, D. Prioritized Experience Replay. arXiv e-prints, art. arXiv: 11111151195222 November 2015.
[52] Schmidhuber, J. Formal theory of creativity, fun, and intrinsic motivation (1990-2010) IEEE Transactions on Autonomous Mental Development, 2(3):230-247, 2010. doi: 10.1109/T
[53] Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal Policy Optimiza- tion Algorithms. arXiv e-prints, art. 1xXiv:1777006344 July 2017.
13

---

[54] Sermanet, P. Lynch, C.. Chebotar, Y., Hsu, J., Jang, E., Schaal, S., and Levine, S. Time-Contrastive Networks: Self-Supervised Learning from Video. arXiv e-prints, art. arXiv: 100.068804..88, April 2017.
[55] Sobel, D., Tenenbaum, J., and Gopnik, A. Children's causal inferences from indirect evidence: Backwards blocking and bayesian reasoning in preschoolers. Cognitive Science, 28:303-333 05 2004. doi: 10.
[56] Srinivas, A., Laskin, M., and Abbeel, P. CURL: Contrastive Unsupervised Representations for Reinforcement Learning. arXiv e-prints, art. arXiv:2004.041 13, April 2020.
[57] Sun, C., Yang. W., Martin, J., and Tonegawa, S. Hippocampal neurons represent events as transferable units of experience. Nature Neuroscience, 23, May 2020.
[58] Sutton, R. and Barto, G. Reinforcement learning: introduction. MIT Press, 2018.
[59] Sutton, R. S., Precup, D., and Singh, S. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning, Artificial Intelligence. 11((1)111-2111 1999. ISSN 0004-3702. doi: http:///wooorp//oooooodo00000010 sciencedirect. wwtpeeetttccmmeeeeee www./o/.ciccaaeeee hwmp/science/artt http:///wwwaacccce/////////////0
URL (RLhhtps://www
[60] Thomas, P. S. and Barto, A. G. Motor primitive discovery. In 2012 IEEE International
Conference on Development and Learning and Epigenetic Robotics (ICDL), pp. 1-8, 2012. doi: 10.1109/Devl...2222..66
[61] van den Oord, A., Li, Y., and Vinyals, O. Representation Learning with Contrastive Predictive Coding, arXiv e-prints, art. arXiv:1801 807.03748, July 2.18.
[62] Widrich, M., Hofmarcher, M., Patil, V. P., Bitto-Nemling A., and Hochreiter, S. Modern hopfield networks for return decomposition for delayed rewards. Deep RL Workshop NeurIPS 2021, 2021. URL
[63] Zacks, J., Braver, T., Sheridan, M., Donaldson, D., Snyder, A., Ollinger, J., Buckner, R., and Raichle, M. Human brain activity time-locked to perceptual event boundaries. Nature Neuroscience, pp. 651-5, June 2001.
[64] Zheng, J., Schjetnan, A. G. P., Yebra, M., Gomes, B. A., Clayton, M., Kalia, K., Valiante, T. A., Mamelak, A. N., Kreiman, G., and Rutishauser, U. Neurons detect cognitive boundaries structure episodic memories in humans. Nature Neuroscience, art. 25, March 2022.
14

---

A
Appendix
A.1
Reproducibility
We aim to maximize the reproducibility of this study. A detailed description and full pseudo-code of ConSpec is given in the method section and appendix. The ConSpec code is available at the link: http://ghuub.coomuuucchhssssss/s I/CnnSpe
A.2 Broader Impact
The focus of this work is the introduction of new add-on algorithm in reinforcement learning to help better generalization and long term credit assignment. We believe that there is potential for positive impact as it can help make many real-life problems, which involve long-term credit assignment, more solvable. However, it could be conceivably applied to ethically quueeallyy problems in RL. Users of such methods must be aware when applying this method to their scientific problems.
A.3 Defining success and failure
How best to define successful and unsuccessful episodes? There are three fairly obvious ways:
1. In the domain of sparse reward tasks, the definition of "successful" episode is relatively straightforwardly defined as an episode that received one of the few rewards available, and "failure" episode defined as anything else.
2. In non-sparse reward settings, the notion of a "successful" episode can be defined as the top set of & episodes based on the sum of rewards. As such, the trajectories defined as successful are updated as an agent improves. Therefore, ConSpec uses the reward signal in comparative and relativistic way, giving it the ability to hone in key steps associated with the best performance achieved so far.
3. In goal-conditioned setting, a "successful" episode is simply one which achieves the goal.
In ConSpec we use the first approach, since the focus of this work is, in part, solving tasks with sparse rewards and multiple contingencies leading to them. In two cases we do take the second approach (with Atari games and Continuous control tasks) (9) to demonstrate that this can also work. We have yet to explore the third approach. which we leave for future research.
A.4 Task implementation details
OrangeTrree experiments were implemented using SilicoLabs' Unity (3I) based game engine to create virtual 3D environments, In stage of this task, the agent sees two objects around the room, one of them round or spiky, and one of them flat (e.g. boxes or tables), and the agent can pick up one of the objects (Fig. 2) In stage 2, the agent is again put in waiting room. In stage 3, the agent must obtain oranges from the tree, but to do so, it must stand on a flat object in order to reach the oranges and receive reward. Successful completion of the task resulted in reward of 10. Observations given to the RL agent were 84x84x3 pixels. The tasks were a total of 65 timesteps long, with 15 timesteps in Stage 1, 40 timesteps in Stage 2, and 10 timesteps in Stage 3.
In the multi-key-to -ey----oo tasks the agent must find all the keys to open the doors (Fig. 4b). the agent does not pick up all of the keys and pass through all of the doors, success is impossible. In between each stage for picking up keys/passing through doors, the agent is forced to wait in waiting room for many time steps in between the retrieval of each key. An episode is successful if and only if the agent exits the final door. Observations given to the RL agent were minimalist, at 5x5x3 pixels. Successful completion of the task resulted in reward of 10. Tasks had alternating key retrieval and waiting periods, and the lengths each of these stages are indicated here A.I
Atari tasks were implemented with frameskips of 4, and for computational expediency, and were uniformly terminated at 600 timesteps at which point rewards were tallied. Observations given to the RL agent were 84x84x3 pixels (where the 3 encodes colours RGB). For Delayed Atari games, episodes were terminated after 600 timesteps and cumulative rewards were given only at the end, similar to previous works 3 48).
15

---

Number of timesteps each task stage:
Stage
Baseline
2-key task
3-key task
4-key task
key-to-door task
Retrieve 1st key
10
10
10
10
Waiting period
85
85
85
85
Retrieve 2nd key
15
Waiting period
85
85
Retrieve 3rd key
Waiting period
85
Retrieve 4th key
Waiting period
Exit the final door
10
10
Total:
405
Figure A.1: Protocols for mult-key-to-doo experiments.
Continuous control tasks were run for 250 timesteps. The traditional version of these tasks gives the agent a forward reward at each timestep based on how much it has moved in the forward direction during that timestep. The delayed versions the tasks- taken from previous literature ((4))) were created by modifying the reward function so that the forward reward is not given at each timestep, but only once the agent has successfully moved a certain number of units in the forward direction (called the threshold). On the timestep that the agent passes threshold, it receives the cumulative reward that it would have received for running up to that threshold during all the previous timesteps. Since a fully trained agent Half Cheetah runs approximately three times farther than fully trained agent hops in Hopper and walks in Walker, the thresholds for Hopper and Walker were set at 1, while the threshold for Half Cheetah was set at 3.
A.5 Experimental and architectural details
Training on all experiments was done on an RTX 8000 GPU cluster. All seeds were run with 32GB of memory until completion.
and F buffer sizes: All experiments used S and J buffer sizes of 16. For 3D OrangeTrre experi- ments, although S and F. buffers each held 16 episodes like all the other experiments, each gradient step done resource limitations. In Atari tasks each gradient step was done on random minibatch of 4 episodes sampled from the buffers of 16 for similar reasons.
random minibatch of 8 episodes sampled from this 16 due computational
Minibatch size: all models used standard minibatch size B 16, except for the Atari tasks, which due to computational resource limitations, used minibatch size B 8.
Seeds experiments: All multi-door-to-key experiments averaged over 10 seeds, did all 3D OrangeTrre experiments. The 9 Atari experiments and 3 Continuous control experiments averaged over seeds each. All plots show median and quartile range.
Number of prototypes: Across implementations of ConSpec for various tasks unless otherwise indicated, the number of prototypes used was set as 8, except for ConSpec on the task in Fig, 16 and the implementation of ConSpec on an RMA backbone A.10), both of which used 16, and the Atari tasks, which due to computational resource limitations, used 3 prototypes.
We also tested ConSpec in the k-k task with various numbers of prototypes and find that performance robust across values (Fig. A.7).
Architectural details for models used: For all multi-key-to-door tasks, the learned encoder of underlying RL agent was single layer 2-D convolutional neural network (outchannels = 32, kernel 3) with ReL.U activation, followed by a final linear layer, and 512-unit GRU, as was done in the (28) codebase for 2D key-to-door tasks. The convnet encoder for the ConSpec module was
16

---

identically sized, but separately parameterized so that there would be no interference between the ConSpec module and the underlying RL agent. In the ConSpec module, the nonlinear projection go in ConSpec was 2-layer MIPP with 1010 units in the intermediate layers and the final output layer, and ReI.U activations between layers. Inspired by the contrastive learning literature, where the importance of having large number of negative samples relative to positive samples has been observed (33), the success buffers were filled at a slower rate than the failure buffers: at most new successful trajectories were added to S from each minibatch, while there was limit to the number of new failed trajectories added to from each minibatch.
Continuous control tasks have small observation spaces and therefore did not require convnet encoder. continuous control experiments, agents simply used a two-layer 64-unit MLP with tanh nonlinearity.
Atari tasks and 3D gridworld tasks possess larger observations, the learned encoder used much larger Impala encoder (4) using 3 Impala layers with 16, 32, and 32 channels respectively. For these experiments, the nonlinear projection ge in ConSpec was 2-layer MLP with 100 units in the intermediate layers and the final output, and ReLU activations between layers.
One of the aims of the Atari experiments definition of successes and failures, as elaborated in section A.3 Specifically, successful trajectories were defined the highest cumulatively scoring ones from each minibatch, while comparatively, failed trajectories were randomly chosen from the remainder of the minibatch. To extend ConSpec to this setup, each of ConSpec's ' prototypes took not just the Sike with the top cosine similarity for each trajectory, but rather, averaged the Sikt's 's with the top-k cosine similarities for each trajectory (where k 20).
to apply ConSpec in setting that requires different
For the hyperparameter A, which scales the intrinsic reward (see equation 3 we used a different value in each set of experiments (multi-key-to-door: Atari, and OrangeTTree). Specifically, we used = 0.5 for the policy-invariant version of ConSpec in the
0.2 for the multi-key multi-key-to- key-to-dorr experiment, À = 0.5 for the Atari games, 0.5 for OrangeTrree and À 2 for Continuous control tasks. These values were determined via a logarithmic grid-search that selected for best reward. Other hyperparameters were found via logarithmic grid-search, and were shared across all tasks:
RL agent learning rate: 2 x 10
ConSpec learning rate: 2 x 10-3
or: 0.2
3: 1.0
We also show for the OrangeTrree task, the 4-key task, and Montezuma's 's Revenge task, that in each of these tasks, performance was relatively insensitive to lambda between 0.2 and 0.5 (Fig. A.9).
Models using PPO backbone used the Adam optimizer with optimal PPO hyperparameters based on the repository (33) (values below). The only major change the hyperparameters used was the entropy coefficient used (0.02, rather than 0.01 from the repository, order to encourage exploration necessary in the multi-key-to-door tasks):
learning rate: 2 10-1
Adam epsilon: 10-5
reward discount factor: 0.99
clipping 0.08
value loss coefficient: 0.5
entropy coefficient: 0.02
PrO-backbone models that were engaged with either the multi-door-to-key tasks or the 3D Orange- Tree task had reward normalized to the range [0. 1). PPO-backbone models that were engaged with the 9 Atari tasks underwent reward clipping to the sign of the reward received at each timestep, as standard practice (39).
17

---

Models using an RMA backbone (including the implementation of TVT) were taken directly from the (28) codebase without modification the architecture. They used the Adam optimizer with PPO hyperparameters based on the repository (28). The only major change to the hyperparameters used was the entropy coefficient used (0.1, rather than 0.05 from the repository, in order to encourage exploration necessary in the nulti-key-to-don tasks):
learning rate: 2 x 10
Adam epsilon: x 10-6
agent discount: 0.92
clipping 0.1
reading strength coefficient: 50.
image strength coefficient: 50.
entropy coefficient: 0.1
The implementation for SynthRs used a synthetics returns MLP that was sized analogously as go in ConSpec in the respective experiments (2-layer MIP with ReLU and 1010 or 100 units respectively). This module was put on top of PPO backbone. Its convolutional neural network encoder was identical to the corresponding encoders used in ConSpec. The version of SynthRs from (48) that did not make use bias term b(s²) was used since in our tasks the reward is totally predictive from the current state (e.g. whether the agent exits the final door not), negating the need for linear regression if the bias term is present. Here. SynthRs learning rates were matched to the learning rates used in ConSpec to enable comparable rates of learning (2 x 10 for the underlying PPO agent, and 2 x 10-3 for the synthetics return module). The optimal weight of the SynthRs loss, 3, as defined in Algorithm 1 was determined via further logarithmic grid-search, and B = 10-4 worked well across experiments. Further hyperparameters were as follows, taken amongst the optimal values used the SynthRs paper (48):
state associate alpha: 0.01
* state-associate beta: 1.0
SynthRs was unable to solve either the 3D OrangeTTree task nor the mult-key-to-door tasks in the allotted time (Figure (3.-b.... But replacing the sigmoid with a softmax on the last layer of the gate g improved performance on both tasks (Figure 4 2b), this was the version of SynthRs used in all experiments.
The implementation of CURL was taken from the repository Impala neural network identical the encoder used by PPO and ConSpec in the 3D OrangeTire experiment. We also replace the SAC agent in the original repo with PPO agent. Hyperparameters were as follows:
We replaced the encoder with an
pretransform image size: 96
image size: 84
frame stack:
batch size for curl: 128
curl latent dim: 128
encoder tau: 0.05
framestack: 100000
The implementation of Decision Transformers for the grid-world task was taken from the repository (6) but we added a convolutional neural network encoder that was identical the encoder used the ConSpec experiment. Moreover, its learning rate (2 x 10-4) 4) was matched to the learning rate used ConSpec to enable comparable rates of learning. All other hyperparmeters were taken without
modification from the repository (6) and were:
weight decay: 1 10-4
18

---

dropout probability: 0.1
context length: 20
number of blocks: 3
The implementation of Random Network Distillation was taken from (30) with the standard convolu- tional neural network encoder that was identical to the encoder used in the ConSpec experiment, All other hyperparameters were the default ones:
extrinsic reward clipping: [-1,1]
intrinsic reward clipping: False
observation clipping after normalization: [-5, 5]
The implementation of Decision Transformers for the Atari tasks was taken from (11 with default parameters, which had been already optimized by the authors for Atari games.
The implementation of RIMs for the Atari tasks was taken from (25) with default parameters, which had been already optimized by the authors for Atari games.
Function for encouraging diversity: The function D is used the loss function, equation for encouraging diversity amongst the prototypes. With the exception of Montezuma's 's Revenge (which will be discuss separately, below), the function that was used for all other experiments was D(( iisig)) Eijz cos (Sta: 8jk) where cos (.;') :) is cosine similarity. This term is therefore an orthogonality constraint imposed to encourage the different prototypes to be independent from each other.
Prototype freezing: keeping prototypes stable once they are learned:
There is potential source of instability in learning prototypes with ConSpec. Specifically, once an agent has learned a prototype for a specific critical step and helped shape the policy so that the agent consistently takes this step, the failure buffer will begin to have many examples where this critical step it is not by itself sufficient to achieve success (e.g. if multiple keys are required). Thereafter the initial prototype would begin to diverge from this critical step, and this can lead to instability, To prevent this, updates to the memory buffers for each prototype were terminated upon reaching criterion that this prototype sufficiently differentiated between successes and failures above threshold T. Across all experiments, this criterion was defined as when the average maximum cosine similarity scores among trajectories S- F differed by > T and that the scores in S were themselves > T, all for at least 25 consecutive gradient steps. We found that T = 0.6 worked well in practice, and this value was used through all experiments.
This design is ceuroscience-motivvee
because semantic knowledge is often defined by exemplars. The frozen set of episodes for prototype i serve as the examplars that define prototype In an experiment, we ablated this prototype freezing mechanism, and interestingly found that even without prototype freezing, ConSpec performs well in the presence multiple contingencies (Fig. A.12).
Recruitment of prototypes:
In most experiments, all prototypes were made immediately available to the agent as the simplest case. However, we also tested recruiting prototypes only as needed in the Montezuma's Revenge and Continuous control experiments. In these experiments, only 3 prototypes were made immediately available to train on. Only when the existing prototypes differentiated success from failure above threshold T, would a new prototype be recruited to be newly trained. Therefore, this was a scheme by which the prototypes would only be recruited as needed.
Montezuma's 's Revenge: The purpose of the Montezuma's 's Revenge experiments was to study whether ConSpec can be used to help Montezuma's 's Revenge, even when it put atop PPO, despite the lack of any further dedicated exploration algorithm. Montezuma's 's Revenge required greater task time (necessary to encourage as much exploration as possible in the environment) and was terminated at longer 1 1200 timesteps. To study this game, the learned encoder was an Impala encoder [1) using 3 Impala layers with 16, 32, and 32 channels respectively, and 256-unit GRU. For these experiments, the nonlinear projection go in ConSpec was 2-layer MIP with 100 units in the intermediate layers and the final output, and ReLU activations between layers. The model used prototypes and X = 0.2. The function D that was used was D((suu Sesssshsse h<ssH) H (L1(E, (2, sue) where (-) is entropy and
19

---

L1(-) is L1 normalization. This term therefore encourages the different prototypes to hone in on the trajectory as uniformly from each other as possible, thereby encouraging prototype diversity. Successes in Montezuma's Revenge that recurrent states stored in the S buffer are rarely updated (especially initially). To alleviate this, we simply replayed randomly chosen success trajectory to the agent every training iteration, to enable it to recalculate its stored recurrent states in S buffer. The training of the ConSpec + PPO model did not initiate until least 2 successes had been spuriously experienced. The minibatch size used in these experiments was B due to limitations in computational resources, and memory freezing was not done as this would have added more memory buffers due to limitations in computational resources. The Montezuma's ' Revenge experiment Fig. (6) used 10 seeds and excluded other seeds that were prematurely terminated by the Pytorch autograd during the learning of the task.
especially sparse especially to the novice agent, which means
Reward
2000
4000
6000
8000
2000
4000 of Ninbatchhs
600D
8000
Minibatches
ConSpec
SynthRs PPO
Figure A.2: Baseline single key-to-door task. (a) Performance of ConSpec on PPO backbone vs SynthRs on PPO backbone vs PPO. (b) Performance of ConSpec on RMA backbone TVT vs RMA.
10
Reward
protottpps prototypes
0
0
4000
8000
Figure A.3: Performance of ConSpec when the number of prototypes is varied (40 8) does change, in the vanilla single key-to-door task.
10
Reward
RUDDER
4000
B000
of Minibatches
Figure A.4: Performance of RUDDER the baseline single key-to-door task.
20

---

RND
Reward
01
0
4000
8000
# of Minibatches
Figure A.5: Performance of RND the baseline single key-to-door task.
Reward
D-Trnddommr
of betches
Figure A.6: ConSpec is more robust than Decision Transformers on vanilla key-to-door because of its objective in learning states rather than predicting actions. Decision transformers is able solve the key-to-door task unlike ConSpec but is able to solve it if actions are constrained in their freedom.
Reward
16 prototypes
12 prototypes prototypes
prototypes
5000
10000
Minibatches
Figure A.7: ConSpec performance was relatively insensitive to changes in the number of prototypes in the 4-key task.
ConSpec: 4-lap task with prototypes
Reward
4000
8000
Minibatches
Figure A.8: ConSpec's performance on the 4-key task with only prototypes: even having fewer than necessary prototypes can surprisingly often catch enough eritical steps still solve the task.
OrangeTree
ke-k-ttooooor task
Montezuma's Revenge
10
10
400
Reward
Reward
Reward
200
0
500
1000
4000
8000
0
1000 2000 3000 of Minibatches
# of Minibatches
of Minibatches
lambda 0.5
lambda 0.5
lambda 0.5
lambda 0.2
lambda 0.2
lambda 0.2
Figure A.9: ConSpec performance was relatively insensitive when À was changed from 0.2 to 0.5 on (a) the 3D OrangeTree task, (b) the 4-key task, and (c) Montezuma's Revenge. 5 seeds each condition.
21

---

10
Reward
4 keys
keys
2 keys
4000
8000
of Minibatches
Performance of ConSpec implemented RMA rather than PPO backbone, the multi-key-to tooo-oo-ooor tasks. As shown, ConSpec was successfully able to rapidly learn and converge in each multi-key-to-dee task, illustrating its flexibility in credit assignment, regardless the underlying RL agent.
Figure A010:
10
Reward
keys
keys
keys
4000
8000
Minibatches
Figure A.11: Performance of ConSpec with the ioviry-invariant version of the intrinsic reward (equation B) on the multi-key-to-donn key-to-door tasks. As shown, ConSpec was successfully able to rapidly learn each task.
10 4
Reward
ConSpec prototypes
ConSpec prototypes frozen
4000
8000
# of Minibatches
Figure A.12: Ablation of the prototype freezing in ConSpec causes no learning deficits in the multikey-to-door task.
ConSpec with stop gradlent
10
Reward
keys
4000
9000
Minibatches
Figure 1.13: ConSpec could not solve the 4-key task with stop-gradient on the prototypes.
RAD
10
Reward
500
1000
Minibatches
Figure A. 14: Performance of RAD the 3D Orange'Tree task.
b
104
SynReg, SynRee, keys SynReg,
Reward
Reward
SynthRs
2000
4000
6000
8000
10000
0
100
200
300
400
500
600
700
Minibatches
Minibatches
Figure A.15: Performance of SynthRs with sigmoid on (a) the multi-key-to-doom and (b) Orange tasks.
22

---

Conjunctive Key Task
Key
Key
Key
Key
Door
T
Waiting
P
Waiting
Waiting
Waiting
ConSpec
TVT
RUDDER
SynthRts
RND
10
leys
10
10
keys
Reward
leys
keys
keys
leys
beys
keys
keys
keys
0
4000
8000
4000
8000
0
4000
8000
4000 8000
4000
8000
Minibatches
Figure A.16: ConSpec on another of multi-key-to-dee tasks with a different, conjunctive structure, Multiple keys had to be learned to be picked up, but picking up each subsequent key did require previous keys to have been successfully picked up, unlike the paper-acceptance inspired tasks in Fig. (a) Protocol for this conjunctive 4 keys task. (b) ConSpec rapidly learns these conjunctive pulti-key-to-door tasks, whereas (c-f) TVT, RUDDER, SynthRs, and RND performances collapse again.
BattleZone
Boxing
Breakout
Reward
60-
60-
70-
20-
2000
4020
2030
4200
2200
4000
Iceltockey
MsPocman
Pong
Revard
1200
800
13.0
2000
4000
Minberchee
Robotank
Spacefnvaders
Roward
400
PPO ConSowe Po Inverians
200
2020
<:00
2000
4000
of Mheberccees
9. Mindbatches
Figure A.17: PPO vs PPO+ConSpec PFO+ConSpec-policy invariant intrinsic reward all learn to perform similarly well on Atari Gym tasks.
RIMs on Breakout
DTransformer on Breakout
Reward
100
1000
2000
3000
1000
2000
Minibatches
Minibatches
Figure A818: Positive controls: Decision transformers and RIMs succeed in learning Atari Breakout,
unlike Montezuma's 's Revenge in Fig.
23


Reviews:
[{"review_id": "hL4H4xGCmL", "reviewer": "Reviewer_t9VM", "summary": "This paper introduces ConSpec, a reinforcement learning (RL) algorithm designed to identify critical steps and improve performance in continuous control tasks. ConSpec utilizes contrastive learning to learn prototypes of critical steps and employs a contrastive loss to differentiate successful and failed experiences. It addresses the challenges of long-term credit assignment and generalization in RL tasks.", "questions": "- ConSpec achieved a return of only 400 in the Montezuma's Revenge task compared to RND, which reached a return of 7500 in its original paper. It appears that ConSpec is not as effective as RND in this regard. Both the multi-key room environment and Montezuma's Revenge involve similar logic, but the former is much simpler. So why does ConSpec outperform RND in Fig.4?\n- It appears that ConSpec has significantly higher variance than the baseline algorithm in all tasks. What could be the reason for this? Further, prototype learning is crucial and can every seed match the key states?\n- What does \"softmaxed over t\" mean in line 152? Why is it necessary to introduce softmax operation?\n- How is success and failure defined in Atari and MuJoCo tasks? I think this is important, but the article lacks details in this aspect.", "limitations": "As presented in Section 5, the number of prototypes is task-specific and definations of  success and failures need human design.  Furthermore,  The proposed method introduces hyperparameters, such as the Diversity measure and the hyperparameter $\\lambda$, that require careful tuning. Furthermore, the algorithm exhibits a significant variance in its actual performance.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This article presents an interesting idea of learning to match key states in a task through contrastive learning. The writing of this paper is clear and Figure 1 is well-drawn, making it easy to quickly grasp the details of ConSpec. And the experimental results effectively demonstrate that the learned prototypes indeed match the key states in 3D Orange-Tree task and gridworld.", "weaknesses": "The cosine similarity measures the similarity between the prototype and the hidden state, both of which are learnable vectors. However, optimizing the contrastive learning loss with updates to both vectors may lead to extremely unstable training. In related literature on representation learning, it is common to use the stop gradient approach and optimize only one of the learnable vectors. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for the feedback.\n\n**Contrastive instability with no stop gradient (SG):** To our knowledge, in self-supervised learning the SG is used (e.g. BYOL Grill et al 2020, or SimSiam Chen et al 2021) to prevent representational collapse (rather than learning instability), and only necessary when there are no negative samples (e.g. SimCLR has no SG). Since ConSpec uses positive and negative samples, there shouldn’t reasonably be SG related learning issues here. Nevertheless, **we did as the Reviewer suggested, stopping gradients to the prototype vectors, and found that it could not solve the 4-key task (Rebuttal Fig. 3), unlike the original ConSpec (Fig. 4c)**. This was likely since SGs cause the prototype vectors to remain in their initial random state, and so, do a poor job of anchoring the critical step recognition process.  We believe that learning both vectors is the correct strategy, and don’t find that it introduces any instability. \n\n**ConSpec v. RND:** The Reviewer is asking why ConSpec vastly outperforms RND, a classic exploration algorithm, on the multikey tasks and not on Montezuma’s Revenge. **This is a good question and reflects a deeper issue**. To us, our results indicate that there is not one but two axes of “difficulty”, and that it is incorrect to consider Montezuma’s Revenge to be more difficult in all respects. Montezuma’s Revenge is more challenging from an exploration stand-point and RND provides sophisticated exploration. On the other hand, RND performance collapses in the multikey tasks. In the  multikey tasks, the pixel difference between “success” and “failure” frames is often very small, and RND struggles to identify these differences. **On the contrary, our contrastive loss is designed to hone in on even very minute differences that distinguish success from failure (figure 4b)**, which is why ConSpec is superior to RND in the multikey tasks, **elegantly illustrating the benefits of ConSpec.**\n\nGiven these two distinct strengths, we wish to highlight that in principle ConSpec could be combined with more powerful exploration algorithms, and so, one should be able to obtain the best of both worlds. \n\n**Variance:** We point out that ConSpec's performance across most tasks display **very little variance at** ***convergence***. This was true in 3D OrangeTree(Fig. 2b), the multi-key task (Fig. 4c), and for delayed Atari (Fig. 5). There does appear to be variance during training, but it seems to be the amount of time it takes for ConSpec to converge to the max reward that varies. In Montezuma's Revenge, variance likely comes from 1) stochasticity in the amount of time for the prototypes to independently hone in on different critical states, and 2) the time it takes to get spurious successes in the first place, just as the Reviewer suspected. But **to us this is benign variance as the seeds do eventually learn, and their performances vastly exceed other baseline algorithms**. Variance is still a major issue in the RL field at large (e.g. Bjorck et al ICLR 2022). \n\n**Softmaxing:** By “softmaxed over t” we mean that we are applying the softmax operation to the scores across time-steps in the trajectory. This was a design choice to sparsify the sequence of cosine scores for each trajectory. Our goal was to compare pairs of prototypes and force their peak cosine scores to be different in time. The softmax was one way to do this, but any equivalent method would do equally well. We will clarify this in a revised manuscript. \n\n**Definition of success/failure:**  In dense reward tasks like Atari and continuous control (Appendix A.3) successes are the top-k highest rewarded trajectories seen so far, while failures were the average trajectories in mini-batches. In RL tasks generally, since there usually is a reward signal, **adaptive thresholding can be used to turn any reward signal into an indicator function of success vs failure**. In this case, ConSpec tries to pinpoint the differences in critical steps between the successes from failures. \n\nTo clarify, in all the experiments in the paper, we made use of only two simple definitions as described in Appendix A.3 (see our general response). We can move this information to the main manuscript if necessary. \n\n**Hyperparameters:** The Reviewer pointed out there are several hyperparameters ConSpec introduces. Luckily, for each hyperparameter, we either found a fixed value that worked across all tasks or a degree of robustness across a range of values: \n1. Number of prototypes: we tested ConSpec with various numbers of prototypes and find that performance is robust across values (Rebuttal Fig. 1).\n2. Diversity measure: Since this hyperparameter acts on cosine similarity scores, we know that it is necessarily bounded between 0 and 1. Moreover, since the point is simply to encourage some non-trivial diversity, we used common sense to select this value, and so, had a fixed value of 0.2 across all experiments. ConSpec learned well across all experiments, speaking to the basic robustness of this parameter choice.\n3.  Lambda: there is a typo in equations (2) and (3) of the manuscript which we apologize for:  $\\tilde{r_{kt}}$ should be equal to $\\lambda \\cdot R_{task} \\cdot \\sum...$ where $R_{task}$ is the nonzero reward per step in the task, under which $\\lambda$’s range is bounded between 0 and 1. We found that performance was relatively insensitive to lambda between 0.2 and 0.5, as shown by new experiments across tasks (Rebuttal Fig. 2) with robust performance.  \n\nAltogether, ConSpec performance was quite stable for these hyperparameters and did not require careful tuning per task.\n\n.\n\nIn sum, this Reviewer brought up a number of issues that we believe actually highlight the robust strengths of ConSpec rather than detract from it. We hope that our responses here, including the new experiments (Rebuttal Fig. 1-5) address the original concerns enough to raise their score, and we are happy to address any insufficiencies.\n"}, {"review_id": "Xyn6dGu1yb", "reviewer": "Reviewer_LNh7", "summary": "Proposes an auxiliary reward module to be used in RL algorithms, that learns features (‘prototypes’) of critical states in successful trajectories. For new observations, the method then uses cosine similarity to the learned features as a reward bonus. The method is evaluated on a unity-based env, grid-worlds, versions of gym,atari envs with delayed rewards, and Montezuma’s revenge.", "questions": "Please address the questions in weakness #1.", "limitations": "Sufficiently addressed", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Effective exploration bonus\n\nThe idea of learning invariant features across successes, and using these as a source of reward does seem to give better exploration performance, from the experiments. The Montezuma’s revenge experiments (Fig,6) are particularly compelling - the baselines PPO, RIMS (which also uses a set of discrete slot-based learned features ) and Decision Transformer all fail to obtain any reward. By creating an explicit division between success and failure episodes, conspec can then learn features that match states present in the successes, but not the failures, even from a very small number of successful trajectories (There might be other, simpler ways to get this effect however, see weakness #1). The ability of con spec to find important states critical for the task is also investigated by the authors in the simpler unity-based env, where they also visualize states closest to the learned prototypes.  \n\n2.   More expressive set for bottleneck states\n\nInstead of learning an explicit set of states which are important (like sub-goals) as has been previously studied, this paper instead captures the notion of ‘critical states’ using learned prototypes. The advantage of this is it can flexibly capture a large set of very different states, all of which are critical. This is also beneficial because it enables zero-shot generalization in new environments (section 4.2)\n\n3. Clarity, presentation\n\nThe paper is well motivated, written clearly, and the main idea for the algorithm is presented clearly.", "weaknesses": "1. Are the prototypes actually required?\n\nLearning from data in successes that aren’t present in failures should lead to better performance, but the importance of doing this through learning prototype features is unclear. As a simple baseline, consider training a policy on only the successful set (using behavior cloning). Does this provide similar performance to con spec on Montezuma’s revenge? Is trying to capture a notion of ‘critical states’ required to learn better policies ? Can you run Decision Transformer where for each successful trajectory, every transition is labelled with a reward of 1, and for every failure trajectory, every transition is labelled with a reward of 0 ?\n\n 2. Success/failure definition\n\nThe method relied crucially on the quality of the learned prototypes, which in turn depends on the success and failure datasets. It might not always be possible to divide up trajectories into 2 classes in this manner, in a lot of tasks performance keeps improving over time and a ‘successful’ trajectory at the beginning of training is very different from one from a converged policy. The authors do discuss this (appendix A.3), but the definition used in this paper for a successful trajectory is - ‘an episode that received one of the few rewards available, and a failure is defined as anything else’. For agents to keep learning and improving from data the notion of a success should necessarily change with time (eg - maximize the reward instead of just getting some reward). \n\n3. Delayed reward envs \n\nA good portion of the experiments are conducted on familiar gym, Atari envs but with a modification where the rewards are delayed. The significance of these experiments is unclear, since the delayed reward setting for these envs is not standard and widespread.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their feedback.\n\n**Are prototypes required?:**\nThe reviewer raises an interesting question: could a system learn critical states implicitly, without resorting to the use of prototypes? As an example, they note that in theory one could use Decision Transformers (DT) and learn from the difference between successful and failed trajectories. Indeed, in Montezuma's Revenge, we have already fed the DT successful trajectories (i.e. trajectories that obtained reward) from an a priori set, as well as failures. In-line with the Reviewer’s recommendation we labeled the successes as 1’s and failure as 0’s. This is the result displayed in Fig. 6d and we found that performance, despite having access to success trajectories, was still zero. The reason is that these successful trajectories used for training did not come from a curated expert policy; rather, they were \"spurious successes\", random-policy trajectories that happened to get reward. This meant that most of the actions taken in these trajectories were actually not the correct actions, but a few were, enough to get reward. As such, behavioural cloning led the DT to learn a lot of terrible policies. Of course, behavioural cloning of curated expert policies with DTs works well. \n\nThis brings up a critical but perhaps subtle point: unlike DTs without prototypes, **ConSpec is able to learn from spurious successes because the prototypes ignore all the states other than those that were critical for distinguishing success from failure. This is why ConSpec’s strategy of honing in on invariant intermediate states and ignoring all the rest of the noise is an important advance, and it is the reason why prototypes are necessary**. Put another way, the secret to the power of contrastively learned prototypes is that their limited capacity (each prototype is a single vector) is essentially a way to force the system to ignore other information that is not relevant to distinguishing successes and failures. \n\nAs our experiments demonstrate, this makes ConSpec, courtesy of its prototypes, much better at learning as a beginner from scratch than DTs, which are designed to learn from experts. \n\n.\n\n**Definition of success/failure:**\nHappily, to clarify, the Reviewer’s recommendation was exactly what we had done. Altogether, we used only two simple definitions of success/failure in all of our tasks as outlined in Appendix A3. One of those definitions was “an episode that received one of the few rewards available”, as the Reviewer commented. But this pertained only to the binary reward setting -- e.g. where there is only 1 reward in the trajectory. \n\nFor dense reward tasks like Atari and Mujoco, successes were defined (Section A.3 line 551-552 in the Appendix) as the top-k highest rewarded trajectories encountered so far. Failures were the median rewarded trajectories in the mini-batches. In this case, success and failure are defined relative to trajectories encountered so far, and as such, the definition of success changes over time as the agent improves - exactly covering the situation the Reviewer describes. **Hence, we agree with the Reviewer and have been doing exactly as the Reviewer recommends all along.** Given this point, we could have been more clear about the fact that the definition of success can change, so we will make this clear in a revised manuscript.\n\nMost interestingly, ConSpec can be a perfect inductive bias to learn from this changing reward signal! ConSpec is designed to pinpoint the relative differences between success and failure trajectories. And **we consider this relativistic design to reflect an important and subtle learning signal that ConSpec exploits but which other RL algorithms do not.**\n\n\n\n.\n\n**Delayed Atari:** Delayed Atari has been used in classic papers like RUDDER, and in more recent papers like InferNet (Ausin et al, 2021). More generally, delayed reward modifications of various regular environments is quite common, including in Synthetic Returns (2021) and in the episodic RL literature (e.g. Off Policy RL with delayed rewards, Han et al 2022), so we do not consider our modified Atari environments out of place. \n\nMoreover, we also apply ConSpec to the standard (unmodified) versions of Atari. Besides Montezuma’s Revenge, we also did so for 8 other Atari games (shown in appendix Fig. A.12). **The significance of our experiments on the delayed Atari is that they take a classic RL benchmark and make it more challenging vis-a-vis long-term credit assignment** (which is why standard RL algorithms such as PPO cannot solve the delayed versions but could solve most of the easier, unmodified Atari games). \n\nUltimately, our use of delayed Atari was but one of many tasks to which we applied ConSpec, including grid world tasks, 3D pixel navigation tasks, regular Atari like Montezuma’s Revenge, and delayed continuous control tasks. These tasks differ from each other in the nature of their observations, their task features, exploration, and goal requirements. The sum total of all these tasks is to showcase the strong capability of ConSpec to robustly handle a wide variety of situations requiring long term credit assignment and/or generalization. \n\n.\n\nAltogether, this Reviewer brought up a number of issues that have helped us to highlight why ConSpec is an important contribution and to clarify our approach. We feel that we have even demonstrated that in some cases ConSpec is doing exactly as the Reviewer had called for. In light of these matters, we hope the Reviewer will raise their score to reflect the novelty and utility of the ConSpec approach, the strength of the varied experiments we provide, as well as the good ideas that we and they jointly thought of.\n"}, {"review_id": "OnhNdJGlXO", "reviewer": "Reviewer_fj2y", "summary": "The paper noticed that in real-world MDP, success is often contingent upon a small set of steps. While Bellman equation can theoretically do credit assignment over long-horizon, reward is hard to propagate under Bellman-based methods in practice. The authors therefore propose a novel algorithm that uses contrastive learning to identify critical states that final success relies on. The method uses a memory like system that can give agents intrinsic reward during training. The paper then evaluates the proposed method on a wide variety of domains and shows performance gain when the proposed method is added to RL algorithms.\n", "questions": "I am wondering whether adding the intrinsic reward can degrade the performance of RL algorithms on common environments (aka, those environments where the final success does not depend on just a few critical steps). This shall be justified the experiments.\n\nWhen the observation is partial, is the proposed method still reasonable?\n\n", "limitations": "1. The method requires privileged information about success of an episode. \n2. I cannot see how the method can be applied to RL that has partial observation that requires recurrent policies.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is based on an interesting and important insight about long-horizon credit assignment and reward learning. The proposed method is designed to explicitly improve long-term credit assignment and have shown empirical success in the evaluation. \n\nThe writing and figures are clear. The paper is easy to follow.\n\nThe evaluation covers a wide variety of RL tasks are benchmarked to back the claim of the paper. ", "weaknesses": "1. The method assumes additional access to a \"success\" indicator at the end of episode. While this is commonly obtainable in gym environments, this doesn't fit into the general MDP setting and thus might limit when the algorithm can be applied. \n\n2. The assumption about access to \"success\" seems privileged compared to baselines. I am wondering they will catch up with the performance of the proposed method when a success bonus is added.\n\n3. The evaluation has #mini batches / # gradient steps as x-axis, unlike the environment steps in common RL benchmarks. I am wondering why this is the case. If this is necessary, I'd like to see convincing justifications. \n\n4. The proposed method relies on a memory system, which may hurt generalization and might have problem when scaling up.\n\n5. CURL+PPO doesn't seem to be a strong baseline to ablate in figure 3. I hope the authors could benchmark against RAD[https://arxiv.org/abs/2004.14990], a much stronger baseline in pixel space. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their questions and feedback.\n\n**Success/failure:**  As discussed in section A.3, we used only 2 simple definitions of success across tasks, and in each case, tie “success” to reward in a natural way. In the first definition, success is based on whether a reward is achieved at all, which works well for sparse reward environments. But, this does not work for dense rewards. Thus, the second definition defines success as simply being the top-k highest rewarded trajectories encountered so far. In other words, there is no fixed definition of success necessary beyond simply being the top level of reward achieved so far. **In this way, success is not privileged information; it is a direct function of the reward observable to the agent, and works with any MDP.**\n\nMoreover, we believe that exploiting the reward signal in this “indicator” way is a special strength of ConSpec. To illustrate, consider dense reward tasks like Atari, where we used the second definition for successes. In this case, ConSpec tries to pinpoint the differences in critical steps between the highest rewarded trajectories and the trajectories with average rewards. **As such, ConSpec uses the reward signal in a comparative and relativistic way, giving it the ability to hone in on what it can improve in each mini-batch, which is a very efficient approach.**\n\n.\n\n**Intrinsic rewards potential degradation of RL performance:** This is a good question and one that we wondered about too. In all of the tasks we have tested to-date, whether they are tasks that require long term credit assignment or not, performance is either drastically improved or unchanged with ConSpec, never degraded. \nWe are confident that this pattern will hold across tasks. This is because degradation of RL performance is in general mitigated by adopting the “potential” form of intrinsic reward (Andrew Ng et al 1999) which provably does not alter the optimal policy of the original task. We have implemented this for ConSpec already, and tested it on key-to-door tasks as well as Atari Appendix Fig. A.8, A.12). \nBut interestingly, even without the “potential” form of intrinsic reward  ConSpec still works well and has never led to degradation on any of the tasks we tested (Fig. 2,4,5,6,7). **We speculate that the reason why performance degradations did not occur with ConSpec is because ConSpec’s intrinsic reward is designed to be small and sparse, and so its unintended negative effects are mitigated.** Nonetheless, as noted, the use of a potential function for calculating the intrinsic reward provably deals with this problem.\n\n.\n\n**Partially observable observations:** Although partially observable environments were not a main focus of the paper, we do have several examples of partially observable environments here. The 3D OrangeTree task was one, since the pixel inputs are by no means a fully observable state. Amongst the Atari tasks, Montezuma’s Revenge and delayed AirRaid were both partially observable, and with these tasks the PPO policy was recursive, and ConSpec was able to give drastic improvements over baselines, so ConSpec was not hindered in any way by partial observability nor recurrent policies.\nMoreover, ConSpec is a plug-and-play module, and is separate from the policy itself. And ConSpec can be considered a representation learning algorithm in some ways (with a novel contrastive loss applied to RL). It discovers promising features separating successes from failures, and it is reasonable to expect it to be agnostic to architectural choices or environment characteristics.\n \n.\n\n**Plotting:** The Reviewer is correct to bring up the plotting conventions. The reason we plotted with # of gradient updates/minibatches is that the ConSpec module is updated based on complete trajectories, so we wanted a metric that would allow us to not only compare baselines, but also, compare across tasks.\n\n.\n\n**Memory system:**  - This is a great point. Scaling to pixel level tasks like Atari games and 3D navigation was not a problem, but testing other larger scale tasks was considered out of scope for the current study. Nevertheless, inspired by the Reviewer, we do have several experiments that begin trying to make ConSpec more compact. To start, the Montezuma experiments don’t fill up the success memory buffers before training starts (since successes are rare) and still works, showing that saving on memory does not hurt. Altogether, the total memory should scale with $~O(k)$ where $k$ is the number of prototypes, which should scale with the number of sparse critical steps, a manageable number.\n\nBut now, we have a new experiment testing learning in ConSpec when the number of prototypes is less than required (e.g. 3 prototypes in the 4-key task) --- in order to study further memory reductions (Rebuttal Fig. 4). Even having fewer than necessary prototypes can often be enough to still solve the task (i.e. catching any critical step helps the agent)! \n\n.\n\n**RAD baseline:** We thank the Reviewer for this suggestion and have done this baseline now (Rebuttal Fig. 5). RAD was still unable to solve the 3D OrangeTree task, unlike ConSpec. To us, this elegantly demonstrates the special benefits of ConSpec’s version of contrastive learning over other representation learning algorithms applied to RL. ConSpec is not merely representation learning; it is a special application of contrastive learning tailored for long term credit assignment and generalization in RL, and our experiments show that it does this extremely well. \n\n.\n\nAltogether, this Reviewer brought up a number of interesting issues that we found very helpful. Through careful examination of these issues as well as interesting new experiments (e.g. 3-prototypes in 4-keys exp, and RAD), we demonstrate that these issues are very mitigatable. We hope the Reviewer may consider raising their score to reflect the novelty and utility of the ConSpec approach, and the strength of the varied experiments.\n"}, {"review_id": "lGH6gPb0E5", "reviewer": "Reviewer_viiH", "summary": "This paper proposes contrastive introspection (ConSpec), an algorithm for learning a set of prototypes for critical states via the contrastive loss. ConSpec works by delivering intrinsic rewards when the current states match one of the prototypes. This paper also conducted experiments in various environments.", "questions": "Per the major weaknesses:\n1. How are the prototypes $h_i$ actually learned? If the reviewer understands correctly, in line 7 of the abstract, the manuscript says “ConSpec learns a set of prototypes…”. While in Algorithm 1, it seems that the prototypes $h_i$ are given to the algorithm as inputs. Maybe the author can clarify why this inconsistency in learning the prototypes happens?\n2. How are the $h_i$ learned/chosen in each experiment? The reviewer has looked into the detail of the experiments in the appendix, but cannot clearly understand how the presented experiments actually utilize the $h_i$. It would be nice that the authors can provide more details of all the $h_i$ in all the present experiments (Sec. 4.1-4.5).  \n", "limitations": "See questions and weaknesses.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The intuition of learning the critical states is natural and easy to follow. The experimental results in this paper look solid and promising.", "weaknesses": "Despite the empirical performance, the reviewer finds the ConSpec algorithm itself hard to follow.\n\nThe largest weakness is: the insufficient discussion on how the prototypes $h_i$ are learned. Hence, the reviewer cannot understand the detailed on how $h_i$ are used (see detailed in Questions). \n\nBesides insufficient discussion on the prototypes, some minor issues are: (1) the title in the pdf (Contrastive Introspection:. ..) seems to mismatch with the one appear in openreview (ConSpec: …). (2) The font of citations appears to be confusing. E.g., from line 19-20 in the introduction, the manuscript uses (number) to address some key points, and the citation also appears as (number) – it would be nice if the citation can be changed to something that is not (number; number).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their feedback.\n\n> *some minor issues are: (1) the title in the pdf (Contrastive Introspection:. ..) seems to mismatch with the one appear in openreview (ConSpec: …). (2) The font of citations appears to be confusing. E.g., from line 19-20 in the introduction, the manuscript uses (number)*\n\nThank you for bringing our attention to these issues. We will address them in turn during revision; we will correct the title, and change the numbering of bullet points to be (i-iv) rather than (1-4) to minimize conflict with the numbered citations (Unfortunately, space constraints prevent us from altering the citations to an author list). \n\n**Main questions:**\nTo answer the Reviewer's main question, the prototypes are initialized as random vectors, and these are provided as initial parameters, per Algorithm 1. However, these initial random vectors are updated based on the gradient of the contrastive loss - they are identical to any other parameter in this way. More specifically, each $h_i$ is a vector of parameters that are updated by the algorithm at step 8 of Algorithm 1. This is stated in section 3.1, 3.2 and 3.3. \n\nWe believe the Reviewer's misunderstanding was due to our listing the prototypes as inputs in Algorithm 1. We see this ambiguity now, so thank you for raising this. But, we note that the prototypes are included in a set of parameters, and this same list includes the other parameters of the model (e.g. the synaptic weights $W$, $\\theta$, and $\\phi$). Thus, these prototypes are only given as inputs to the algorithm at the start in the same way that all the other randomly initialized parameters are. To clarify this point and avoid any other confusion we will separate out the parameters from the other inputs in Algorithm 1 in revision.\n\nFurthermore, we hope this resolves the Reviewer’s second question about how we choose the prototypes: we do *not* choose them. They should only be understood as parameters that are updated via gradient descent. \n\n.\n\nConceptually, **ConSpec uses these prototypes in a very novel way, with a new contrastive loss, providing a fresh solution to both difficult problems of long-term credit assignment and generalization in RL, as we demonstrate through a variety of different task situations.**\n\nWe hope that with this clarification on how the ConSpec procedure works, the reviewer will enjoy the merits of this novel algorithm and their score will be updated appropriately. \n"}]

Author rebuttal:
**Message to all reviewers:**
Thank you for the comments, questions, and suggestions. In this global response we will address those issues highlighted by multiple reviewers. We will respond to individual reviewer comments in the individual responses.

First, we want to summarize and clarify the central, novel contribution of our work. Whereas most contemporary RL algorithms try to model or fit values, returns, or transitions for all the states in the environment, the present manuscript investigates an alternative learning strategy: that of retrospectively honing in on a few critical steps, courtesy of a novel contrastive loss. ConSpec tests the hypothesis that **the ability to recognize  critical steps (and ignore all other states) allows rapid long-term credit assignment and robust generalization.** In support of this hypothesis, we have demonstrated strong performances by ConSpec in a wide variety of different task environments, discrete and continuous tasks, gridworlds, 2D Atari including Montezuma’s Revenge, and 3D navigation. We managed to isolate particular weaknesses of other algorithms for long-term credit assignment including RUDDER, Synthetic Returns, and Temporal Value Transport (TVT) in situations with multiple contingencies, which are common in real life, and which ConSpec overcomes easily thanks to its retrospective credit assignment. Moreover, we even managed to **isolate particular weaknesses** of other RL strategies like RND exploration, and Decision Transformers that ConSpec can overcome, demonstrating the uniqueness of ConSpec as well as its potential complementarity to these other algorithms. 

Thus, our results show that retrospectively honing in on critical steps is a potentially powerful thing to do especially for agents learning from scratch (to learn the world a bit at a time). 

**New experiments**: In the individual Reviewer questions below, the Reviewers brought up a variety of interesting issues that have helped us to think more deeply about this approach. Motivated by these points, we conducted a number of small experiments. We found that all the worries that the Reviewers brought up are easily mitigated by ConSpec. **Our new experiment figures are in the attached PDF, and include:**

- Demonstrating robustness to hyperparameter choices
- A version of ConSpec with stop-gradient on prototypes 
- More compact ConSpec with fewer than necessary prototypes 
- A new baseline RAD (stronger than CURL) 

**Success v. Failure**: Several Reviewers brought up issues related to the definition of successes (Reviewers t9VM, LNh7, fj2y), so we will give a summary here (and we have also answered their individualized questions below). As discussed in Appendix section A.3, we used only 2 simple definitions of success across all of our tasks, and in each case, we tied “success” to reward in a natural way. In the first definition, success is based on whether a reward is achieved at all, which works well for sparse reward environments, but this does not work for dense rewards. Thus, the second definition defines success as the top-k highest rewarded trajectories encountered so far, which means that the trajectories defined as successful are updated as an agent improves. As such, ConSpec uses the reward signal in a comparative and relativistic way, giving it the ability to hone in on key steps associated with the best performance achieved so far, a very efficient approach. This is a special strength of ConSpec. 

In light of our findings, we believe that ConSpec is a **robust and versatile plug-and-play** module that can improve credit assignment and generalization in RL. Moreover, it is a **novel solution to traditional RL problems**. Most notably, it is a novel application of **contrastive** learning in RL, distinct from other contemporary attempts, and it takes a **retrospective** strategy to credit assignment, discovering prototypes of a few critical steps that have the capacity to both help assign credit over long time horizons and generalize, as we show with strong performance in widely varied experimental settings. 

All in all, we believe that ConSpec provides a novel, well-validated approach to improving difficult credit assignment in RL, and would benefit the NeurIPS community in this venue. We hope that the reviewers agree, particularly given our new experiments, and see fit to raise their scores.


Meta-review: The paper provides a unique and interesting approach to representation learning in RL. I went through the reviews and paper. I agree with the reviewers that the paper is a strong submission: Good and clear writing, interesting and novel algorithmic approach, well-designed experiments (thanks to the reviewers who encouraged authors to run a few more ablations/baselines!).
Decision: Accept (poster)