Learning State Representations from Random Deep Action-conditional Predictions

1 Introduction . Providing auxiliary tasks to Deep Reinforcement Learning ( Deep RL ) agents has become an important class of methods for driving the learning of state representations that accelerate learning on a main task . Existing auxiliary tasks have the property that their semantics are fixed and carefully designed by the agent designer . Some notable examples include pixel control , reward prediction , termination prediction , and multi-horizon value prediction ( these are reviewed in more detail below ) . Unlike the prior approaches that require careful design of auxiliary task semantics , we explore here a different approach in which a set of random action-conditional prediction tasks are generated through a rich space of general value functions ( GVFs ) defined by a language of predictions of random features of observations conditioned on a random sequence of actions . Our main , and perhaps surprising , contribution in this work is an empirical finding that auxiliary tasks of learning random GVFs—again , random in both predicted features and actions the predictions are conditioned upon— yield state representations that produce control performance that is competitive with state-of-the-art auxiliary tasks with hand-crafted semantics . We demonstrate this competitiveness in Atari games and DeepMind Lab tasks , comparing to multi-horizon value prediction [ 8 ] , pixel control [ 13 ] , and CURL [ 15 ] as our baseline auxiliary tasks . Note that while we present a reasonable approach to generating the semantics of the random GVFs we employ in our experiments , the specifics of our approach is not by itself a contribution ( and thus not evaluated against other approaches to producing semantics for random GVFs ) , and alternative reasonable approaches for generating random GVFs could do as well . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Additionally , through empirical analyses on illustrative domains we show the benefits of exploiting the richness of GVFs—their temporal depth and action-conditionality . We also provide direct evidence that using random GVFs learns useful representations for the main task through stop-gradient experiments in which the state representations are trained solely via the random-GVF auxiliary tasks without using the usual RL learning with rewards to influence representation learning . We show that , again , surprisingly , these stop-gradient agents outperform the end-to-end-trained actor-critic baseline . 2 Background and Related Work . Horde and PSRs . Auxiliary tasks were formalized and introduced to RL in [ 27 ] through the Horde architecture . Horde is an off-policy learning framework for learning knowledge represented as GVFs from an agent ’ s experience . Our work is related to Horde in the use a rich subspace of GVF predictions but differs in that our interest is in the effect of learning these auxiliary predictions on the main task via shared state representations rather than to show the knowledge captured in these GVFs . Our work is also related to predictive state representations ( PSRs ) [ 17 , 24 ] . PSRs use predictions as state representations whereas our work learns latent state representations from predictions . Recently , in the use of deep neural networks in RL as powerful function approximators , various auxiliary tasks have been proposed to improve the latent state representations of Deep RL agents . We review these auxiliary tasks below . Our work belongs to this family of work in that the auxiliary prediction tasks are used to improve the state representations of Deep RL agents . Auxiliary tasks using predefined GVF targets . UNREAL [ 13 ] uses reward prediction and pixel control and achieved a significant performance improvement in DeepMind Lab but only marginal improvement in Atari . Termination prediction [ 14 ] is shown to be an useful auxiliary task in episodic RL settings . SimCore DRAW [ 9 ] learns a generative model of future observations conditioned on action sequences and uses it as an auxiliary task to shape the agent ’ s belief states in partially observable environments . Fedus et al . [ 8 ] found that simply predicting the returns with multiple different discount factors ( MHVP ) serves as effective auxiliary tasks . MHVP relies on the availability of rewards and thus is different from our work and other unsupervised auxiliary tasks . Information-theoretic auxiliary tasks . Information-theoretic approaches to auxiliary tasks learn representations that are informative about the future trajectory of these representations as the agent interacts with the environment . CPC [ 30 ] , CPC|action [ 10 ] , ST-DIM [ 1 ] , DRIML [ 19 ] , and ATC [ 25 ] apply different forms of temporal contrastive losses to learn predictions in a latent space . CURL [ 15 ] ignores the long-term future and applies a contrastive loss on the stack of consecutive frames to learn good visual representations . PBL [ 11 ] focuses on partially observable environments and introduces a separate target encoder to set the prediction targets . The target encoder is trained to distill the learned state representations . SPR [ 23 ] replaces the target encoder in PBL with a moving average of the state representation function . In addition to being predictive , PI-SAC [ 16 ] also enforces the state representations to be compressed . SPR and PI-SAC focuses on data efficiency and only conducted experiments under low data budgets . In these information-theoretic approaches , the targets are not GVFs and despite some empirical success , none could learn long-term predictions effectively . This is in contrast to GVF-like predictions which can be effectively learned via TD as in our work as well as the work presented above . Theory . A few recent works [ 3 , 7 , 18 ] have studied the optimal representation in RL from a geometric perspective and provided theoretical insights into why predicting GVF-like targets is helpful in learning state representations . Our work is consistent with this theoretical motivation . GVF discovery . Veeriah et al . [ 31 ] used metagradients to discover simple GVFs ( discounted sums of features of observations ) In this work , we show that random choices of features and random but rich GVFs are competitive with the state-of-the art of hand-crafted GVFs as auxiliary tasks . GVF RNNs . Rather than using GVFs as auxiliary tasks , General Value Function Networks ( GVFNs ) [ 22 ] are a new family of recurrent neural networks ( RNNs ) where each dimension of the hidden state is a GVF prediction . GVFNs are trained by TD instead of truncated backprop through time . Our work relates to GVFNs in that both works use GVFs to shape state representations . However , unlike GVFNs , our work uses GVFs as auxiliary tasks and does not enforce any semantics to the state representations . Moreover , our empirical study mainly focuses on the control setting where the agent needs to maximize its long-term cumulative rewards , whereas [ 22 ] mainly focuses on time series modelling tasks and online prediction tasks and demonstrated the superior performance of GVFNs over conventional RNNs when the truncated input sequences are short during training . 3 Method . In this section we first describe the specific GVFs we studied in this work . Then we describe an algorithm for the construction of random GVFs . We finish this section with a description of the agent architecture used in our empirical work . 3.1 GVFs with Interdependent TD Relationships . In this work , we study auxiliary prediction tasks where the semantics of the predictions are defined by a set of GVFs with interdependent TD relationships ( this family of GVFs are often referred as temporal-difference networks in the literature [ 28 , 29 ] ) . The TD relationships among the GVFs can be described by a graph with directional edges , which we call the question network as it defines the semantics of the predictions . Figure 1a shows an example of a question network . The two squares represent two feature nodes and the six circles represent six prediction nodes . Node 1 ( labeled in the circles ) predicts the expected value of feature f1 at the next step . Implicitly , this prediction is conditioned on following the current policy . Node 2 predicts the expected value of node 1 at the next step . Note that we can “ unroll ” the target of node 2 to ground it on the features . In this example , node 2 predicts the expected value of feature f1 after two steps when following the current policy . Node 5 has a self-loop and predicts the expectation of the discounted sum of feature f2 with a discount factor γ . We call node 5 a discounted sum prediction node . Node 3 is labeled by action a . It predicts the expected value of node 1 at the next step given action a is taken at the current step . We say node 3 is conditioned on action a . Similarly , node 4 predicts the same target but is conditioned on action b. Node 6 has two outgoing edges . It predicts the sum ( in general a weighted sum , but in this paper we do not explore the role of these weights and instead fix them to be 1 ) of feature f1 and the value of node 5 , both at the next step . In this case , it is hard to describe the semantic of node 6 ’ s prediction in terms of the features , but we can see that the prediction is still grounded on feature f1 and f2 . Generalising from the example above , a question network with np prediction nodes and nf feature nodes defines np predictions of nf features . We use Np to denote the set of all prediction nodes and Nf to denote the set of all feature nodes . Let W be the adjacency matrix of the question network . Wij denotes the weight on the edge from node i to node j . We define Wij , 0 if there is no edge from node i to node j . Now consider an agent interacting with the environment . At each step t , it receives an observation Ot and takes an action At according to its policy π . Then at the next step it receives an observation Ot+1 . The feature fk ( Ot , At , Ot+1 ) is a scalar function of the transition . The agent makes a prediction ŷi ( O0 , A0 , . . . , Ot ) for each prediction node i based on its history ; this is computed by a neural network in our work . For brevity , we use fkt+1 and ŷ i t to denote fk ( Ot , At , Ot+1 ) and ŷi ( O0 , A0 , . . . , Ot ) respectively . The target for prediction i at step t is denoted by yit . If prediction node i is not conditioned on any action , its target is yit = Eπ [ ∑ j∈Np Wijy j t+1 + ∑ k∈Nf Wikf k t+1 ] otherwise , if it is conditioned on action ai , its target is yit = Eπ [ ∑ j∈Np Wijy j t+1 + ∑ k∈Nf Wikf k t+1|At = ai ] . By the construction of the targets , the agent can learn the prediction ŷit via TD . If i is not conditioned on any action , then ŷit is updated by ŷit ← ∑ j∈Np Wij ŷ j t+1 + ∑ k∈Nf Wikf k t+1 otherwise , if i is conditioned on action ai , then ŷit is updated by ŷit ← ∑ j∈Np Wij ŷ j t+1 + ∑ k∈Nf Wikf k t+1 if At = a i ŷit otherwise In an episodic setting , if the episode terminates at step T , we define yiT , 0 and ŷ i T , 0 for all prediction nodes i . These GVFs represent a broad class of predictions . Many existing auxiliary prediction tasks can be expressed by a question network . Reward prediction [ 13 ] can be represented by a question network with a single feature node representing the reward and a single prediction node predicting the reward . Multi-horizon value prediction [ 8 ] can be represented by a similar question network but with multiple self-loop prediction nodes with different discounts . Termination prediction [ 14 ] can be represented by a question network with a feature node of constant 1 and a self-loop node with discount 1 .

Summary: This paper introduces a new auxiliary task for learning state representations. This is done by predicting general value functions (GVFs), which are predictions of random features of observations conditioned on sequences of actions. The GVFs of interest can be represented with graphs which they call "question networks". The edges in this graph define nodes predicting the value of other nodes in a previous time-step, and many existing prediction tasks in the literature such as reward prediction, multi-horizon value prediction, and termination prediction can be represented using a question network.  The key contribution here is that predicting the value of nodes in a randomly generated question can provide a strong learning signal for learning representations. Several experiments are provided that show that: (1) representations learned by predicting the values of large random question networks (rGVF) outperform baselines like pixel-control and MHVP; (2) rGVF was the only method that learned good representations even when using a stop-gradient to prevent the RL learning signal from shaping the representation; and (3) several ablation tests were done to show the importance of action conditioning and depth.