L EARNING C ONTROLLABLE E LEMENTS O RIENTED R EPRESENTATIONS FOR R EINFORCEMENT L EARNING Anonymous
authors
Paper under double-blind review
A BSTRACT
Deep Reinforcement Learning (deep RL) has been successfully applied to solve various decision-making problems in recent years.
However, the observations in many real-world tasks are often high dimensional and include much task-irrelevant information, limiting the applications of RL algorithms.
To tackle this problem, we propose LCER, a representation learning method that aims to provide RL algorithms with compact and sufﬁcient descriptions of the original observations.
Speciﬁcally, LCER trains representations to retain the controllable elements of the environment, which can reﬂect the action-related environment dynamics and thus are likely to be task-relevant.
We demonstrate the strength of LCER on the
DMControl Suite, proving that it can achieve state-of-the-art performance.
To the best of our knowledge, LCER is the ﬁrst representation learning algorithm that enables the pixel-based SAC to outperform state-based SAC on the DMControl 100K benchmark, showing that the obtained representations can match the oracle descriptions ( i.e. the physical states) of the environment.
I NTRODUCTION
Deep Reinforcement Learning (deep RL) has proven its ability to solve difﬁcult sequential decisionmaking problems such as Dota2 (Berner et al., 2019) and StarCraft (Vinyals et al., 2019).
However, it is still a challenge to apply deep RL to many real-world tasks, because the observations in these tasks are often high dimensional ( e.g. pixels) and include much task-irrelevant noise.
To address these issues, representation learning is introduced in deep RL to provide information-dense descriptions of the original observations, and thus reduce the difﬁculty of solving RL problems.
Many representations learning algorithms are proposed with the design goal to make the representations be predictive of properties of future states.
Such a goal is typically achieved by imposing the representations to reﬂect the dynamics of the environment (Mazoure et al.
; Lee et al., 2020b;
Schwarzer et al., 2020; Gelada et al., 2019).
However, the dynamics of the environment are not always task-relevant( e.g. the noisy and varying background in control tasks).
This can be problematic because those task-irrelevant dynamics will entice the representations into encoding task-irrelevant information( e.g. the noisy background).
Since the capacity of representations is limited, the obtained representations can miss important task-relevant information (Gelada et al., 2019) or encode too much unimportant task-irrelevant information (Zhang et al., 2020).
Therefore, it is essential to learn representations with more task-relevant environmental dynamics.
We envision that action-related dynamics are likely to be more task-relevant because actions are the only feedback from the agent to the environment.
For example, in control tasks and robotics particularly, the agent can complete its task only if it can realize which part of the environment can be inﬂuenced by its actions.
For simplicity, we use controllable elements to denote those environment elements that can be inﬂuenced by the agent’s actions and thus reﬂect the action-related dynamics.
We argue that these controllable elements should be explicitly retained for better representations.
Following the intuition above, we develop an efﬁcient and ﬂexible framework named L earning
C ontrollable E lements oriented R epresentations (LCER) for reinforcement learning to capture the controllable elements of the environment.
Speciﬁcally, we formally deﬁne a metric to measure the amount of controllable elements that representations capture, and then propose a surrogate objective that is derived from this metric and can make the representations encode controllable elements as well as sufﬁcient information.
By emphasizing the importance of the controllable elements, LCER can capture more task-relevant information automatically.
We evaluate our method LCER on the DMControl Suite (Tassa et al., 2018), which has been widely used in recent years to benchmark sample efﬁciency (Kostrikov et al., 2020b; Hafner et al., 2019;
Laskin et al.
; Lee et al., 2020b).
Our experiments show that LCER combined with SAC (Haarnoja et al., 2018) can outperform other representation learning methods on a majority of tasks, and even achieve better performance than the state-based SAC on the DMControl 100K benchmark.
We also evaluate LCER on the Distracting Control Suite (Stone et al., 2021), showing LCER’s ability to ﬁlter out distracting factors.
The implementation of our work is available at https://anonymous.
4open.science/r/LCER-FF  R ELATED W ORK
In this section, we review the common approaches of representation learning in RL brieﬂy.
Predictive representations in RL.
Predictive representations are typically trained to predict the properties of future states.
One commonly used approach is to predict the representation of future states given the current state’s representation and actions.
Such an approach is widely used in modelbased RL algorithms.
The representations are prone to collapse if without further constraints (Gelada et al., 2019), thus a reconstruction loss is usually applied to ensure non-trivial representations.
However, this approach tries to encode everything, unable to ﬁlter out task-irrelevant information.
SPR (Schwarzer et al., 2020) introduces a target encoder to avoid collapse and reconstruction, but its effect lacks theoretical guarantee, and it dose not emphasize the controllable elements as LCER does.
There are some methods that train predictive representations from information-theoretic perspectives, such as (van den Oord et al., 2018;
Hjelm et al., 2019; Anand et al., 2019; Lee et al., 2020b; Mazoure et al.).
These methods do not give a bias explicitly towards what kind of environment’s dynamics should the representations reﬂect, while our method LCER focuses more on the actionrelated dynamics which is more likely to be task-relevant.
Our experiments in Section 5.2 shows that encoding predictive information does not ensure representations of high quality in complex environments, and it is crucial to retain controllable elements explicitly.
There are some prior works that also emphasize controllable elements. Typically, (Zhang et al., 2018; pok, 2016; Badia et al., 2020) learn an inverse model that predicts the action given successive states. However, these approaches typically rely on a reconstructive loss, which will signiﬁcantly harm their performance in complex environments. There are also some works that emphasize the importance of actions in representation learning by introducing various action-related objectives such as policy selectivity(Bengio et al., 2017) and contingency awareness(Bellemare et al., 2012).
However, their methods are often built on heuristics and do not give a clear deﬁnition of controllable elements.
Compared with these works, LCER ﬁrstly introduces a novel measure of controllable elements (Eq.(3)) and then develops an efﬁcient method to optimize it without the aid of reconstruction loss.
Bisimulation-based Representations in RL.
Bisimulation-based representations group those states into a cluster which are indistinguishable w.r.t.
future return, given any action sequence tested.
The indistinguishability of two states is typically measured by the bisimulation metrics, which are ﬁrstly introduced in (Ferns & doina Precup, 2014).
Recently, (Zhang et al., 2020) proposes DBC, a representation learning algorithm based on bisimulation that presents impressive results in distracting environments.
However, these methods rely on a well-deﬁned reward function, which is probably unavailable in real-world tasks.
Besides, they often fail to solve simple tasks efﬁciently such as ones from DMControl Suite (Zhang et al., 2020).
Prior-knowledge-based Representations in RL There are other methods that use prior knowledge to constraint the representations (Jonschkowski & Brock, 2015; Jonschkowski et al., 2017;
Thomas et al., 2018).
For example, (Jonschkowski & Brock, 2015) forces the representations to satisfy multiple priors such as proportionality and repeatability, which generally describe how the optimal representations should look like.
However, such priors need to be specially designed for each environment, and much domain knowledge is required to form good representations.
These drawbacks signiﬁcantly limit the applications of these methods.
P RELIMINARIES 3. N OTATION
We assume the underlying environment is a Markov decision process (MDP), described by the tuple
M = ( S, A, P, R, γ ) , where S is the state space,
A the action space, P : S t × A t × S t +1 → [0 , 1] the transition probability function which determines the distribution of next state given current state and action, and γ ∈ [0 , 1] a discount factor.
Given the current state s ∈ S , an agent choose its action a ∈ A according to a policy function a ∼ π ( ·| s ) .
This action will update the system state to a new state s (cid:48) according to the transition function P , and then a reward r = R ( s, a, s (cid:48) )
∈ R is given to the agent.
The goal of the agent is to maximize the expected cumulative rewards by learning a policy π .
We denote the state at time step t by s t .
We use the upper letter of s ( i.e. S ) to refer to the random variable of state if there is no ambiguity.
We use s t : t + k to present the state sequence from t to t + k , i.e. s t : t + k :=
( s t , s t +1 , ..., s t + k ) .
These notations can be extended to other variables such as actions and rewards.
We use φ : S → Z to denote the embedding function of LCER, which maps s ∈ S to a latent space Z .
3. M UTUAL I NFORMATION
The mutual information(MI) of two random variables X, Y is deﬁned as:
[Equation]
where D KL is the Kullback–Leibler divergence that measures the difference of two distributions.
The conditional mutual information(CMI) I ( X ; Y | Z ) is deﬁned as the expected value of I ( X ; Y ) given the value of Z :
[Equation]
L EARNING C ONTROLLABLE E LEMENTS O RIENTED R EPRESENTATIONS
Given an embedding function: φ : S → Z , we argue that the number of controllable elements that φ contains can be measured by the conditional mutual information(CMI):
[Equation]
Intuitively, I k measures the information that φ ( S t + k ) and A t : t + k − 1 shares given φ ( S t ) .
Note that the non-controllable elements captured by φ do not contribute to I k because they are not related to the actions.
Training representations to maximize Eq.(3) is an efﬁcient way to make the representations capture controllable elements.
However, representations obtained in this way may be too compressive 1 to provide sufﬁcient information for solving tasks in some cases, because not all task-relevant elements are controllable.
For example, in navigation tasks, the position of the goal is task-relevant but not controllable by the agent.
Therefore, we need to utilize Eq.(3) in a more delicate manner. Speciﬁcally, in Section 4.1 we propose a surrogate objective derived from Eq.(3) which additionally provides a mechanism to control the degree of compression.
In Section 4.2, we describe how to optimize this objective, leading to our ﬁnal algorithm LCER.
4. T HE O PTIMIZATION O BJECTIVE FOR LCER
By the chain rule of mutual information, we can break the I k deﬁned in Eq.(3) into the subtraction of two terms: I k ( φ )
=
I ([ φ ( S t ) , A t : t + k − 1 ]; φ ( S t + k )) − I ( φ ( S t ); φ ( S t + k )) .
The ﬁrst term tries to capture predictive information, while the second term serves as a compression term that aims to ﬁlter out action-irrelevant information.
To make φ capture the controllable elements and maintain sufﬁcient information at the same time, we introduce a hyper-parameter β to control the degree of compression:
[Equation]
By choosing different β when training φ to maximize Eq.(4), we can control the degree of compression, making it possible to adapt to different environments by choosing different β .
However, directly maximizing I k ( φ ; β ) is intractable, because maximizing the ﬁrst term and minimizing the second term in Eq.(4) at the same time have an antagonistic effect to each other, thus it is difﬁcult to tune in practice.
In fact, we ﬁnd this will often lead to unstable training process in our early experiments.
To tackle this problem, we substitute I ( φ ( S t ); φ ( S t + k )) with I k ( φ ; β = 0) − I k
[Equation]
φ ;
[Equation] [Equation]
[Equation]
In Eq.(5), we maximize two terms ( J 1 k and J 2 k ) at the same time, and their coefﬁcients are 1 − β and β respectively.
The ﬁrst term J 1 k is closely related to the predictive information encoded by φ .
The second term J 2 k is related to the controllable elements, which will play an important role in distracting environments as shown in Section 5.3.
In practice, we sum over different choices of k , leading to the ﬁnal objective to maximize:
[Equation]
4. H OW TO O PTIMIZE THE O BJECTIVE FOR LCER
To maximize I K ( φ, β ) , we propose to maximize its corresponding lower bound, which is common practice in mutual information estimation.
In this section, we ﬁrst provides the lower bound for J 1 k and J 2 k respectively that are used in this paper, and ﬁnally put them together to get our ﬁnal loss for training φ .
Optimizing J 1 k .
In order to maximize J 1 k ( φ ) , we propose to maximize the InfoNCE (van den Oord et al., 2018) lower bound of J 1 k ( φ ) :
[Equation]
where f is any scalar function, B is the batch size, and s it , a it : t + k − 1 , s it + k is a transition sequence from s it to s it
+ k .
There are other alternatives to InfoNCE such as NWJ (Poole et al., 2019) and
MINE (Belghazi et al., 2018), however InfoNCE can often lead to better representations as shown in (Tschannen et al., 2020).
In order to efﬁciently optimize ˆ J 1 k for all 1 ≤ k ≤ K , we recursively deﬁne ψ it + k to encode the information of ( φ ( s it ) , a it : t + k − 1 ) :
[Equation]
otherwise
[Equation]
and set f to be the bi-linear inner product of ψ it + k and φ ( s it + k ) :
[Equation]
Note that the φ it + k is shared across all ˆ J 1 k , thus we do not need to re-compute it.
The h in Eq.(8) is quite similar to a “prediction” network, but is trained together with φ instead of prediction errors.
Optimizing J 2 k .
In order to maximize J 2 k ( φ ) , we turn to maximizing the corresponding JensenShannon-divergence-based CMI
I JSD :
= I JSD ( φ ( S t + k );
A t : t + k − 1 | φ ( S t )) 2 , because I JSD is proven to be more stable in our experiments, which is also suggested in prior works (Hjelm 2 I JSD ( X ; Y | Z ) = E [Equation] ∼
[Equation]
D JS ( P XY | z (cid:107) P X | z P Y | z ) , where D JS is the Jensen-Shannon divergence et al., 2019; Sanchez et al., 2020).
To maximize I JSD , we introduce a statistics network T k :
S ×
A t : t + k − 1 × S → R and maximize:
[Equation]
where j (1) , .., j ( B ) are obtained by randomly shufﬂing 1 , .., B .
The derivation of Eq.(10) is provided in Appendix B.
In practice, we choose T k ( φ t , a t : t + k − 1 , φ t + k )
=
T ( ψ t + k , φ t + k ) , where ψ t + k is a latent deﬁned in Eq.(8) and T : Z × Z → R is a network shared across different T k .
Final loss for training φ .
Inspired by (He et al., 2020), we propose to utilize a target encoder φ targ to stabilize the training process, whose weights are the exponential moving average(EMA 3 ) of φ .
Speciﬁcally, the embeddings of s it + k ( k ≥ 1) are calculated by φ targ .
All pieces together, we obtain our loss for φ to minimize:
L K ( φ ; h, W, T ) = −
[Equation]
1 + exp ( − T ( ψ it + k , φ targ ( s it + k ))) + log(1 + exp ( T ( ˆ ψ it + k , φ targ ( s it + k )))] , (11) where B is the batch size, h, ψ it + k are deﬁned in Eq.(8), W is deﬁned in Eq.(9), ˆ ψ it + k is a latent deﬁned in the same way with ψ it
+ k
but it encodes φ ( s it ) and a j ( i ) t : t + k − 1 (instead of a it : t + k − 1 ), T is a scalar function, and φ targ is the target encoder.
The parameters of φ, h, W, T are training together by minimizing L K .
In practice, Eq.(11) can be used as an auxiliary loss for any RL algorithms such as SAC (Haarnoja et al., 2018).
E XPERIMENTS
In this section, we ﬁrst compare LCER with other methods in Section 5.1 on the PlaNet benchmark.
In Section 5.2, we then demonstrate LCER’s strength to ﬁlter out distracting factors on the Distracting Control Suite, and show that it is not enough to obtain good representations just by encoding predictive information.
In Section 5.3, we investigate the effect of K and β , which are the most important parameter in LCER.
Finally, we demonstrate LCER can obtain representations of higher quality than the vanilla SAC, as well as its potential to transfer learning in Section 5.4.
In our experiments, LCER is equipped with SAC (Haarnoja et al., 2018).
Most of the parameters in
LCER are the same with PI-SAC (Lee et al., 2020b), including the batch size, learning rate, action repeat, and the network architecture of φ , see Appendix A for more details.
Following (Laskin et al.;
Lee et al., 2020a;b), we report the performance using the true environment steps, which is invariant to the choice of action repeat.
Unless specially speciﬁed, We set K = 5 , β = 0 .
1 and run every experiment using 10 random seeds.
5. E VALUATION ON S AMPLE E FFICIENCY
In this section, We evaluate LCER on the PlaNet benchmark, which consists of six challenging control tasks from the DMControl Suite (Tassa et al., 2018).
The PlaNet benchmark is ﬁrst introduced in (Hafner et al., 2019) and later widely used to benchmark sample efﬁciency in (Kostrikov et al., 2020b; Laskin et al.
; Lee et al., 2020b;a).
We equip LCER with SAC (Haarnoja et al., 2018) and compare LCER to 3 leading representation learning algorithms 4 : CURL (Laskin et al.),SLAC
(Lee et al., 2020a),and PI-SAC (Lee et al., 2020b).
[Table] CURL introduces contrastive learning into RL and actually trains representations to encode information that is invariant under data augmentation.
SLAC efﬁciently combines both representation learning and control into a single objective via Bayesian inference [Table] but it relies on a reconstruction loss.
PI-SAC tries to encode predictive information, while LCER focuses on encoding controllable elements.
All these methods are also built on SAC (Haarnoja et al., 2018).
Following (Laskin et al.; Kostrikov et al., 2020b), we use the version of PI-SAC that performs one gradient update per environment step to ensure a fair comparison.
We also compare LCER with a model-based algorithm PlaNet (Hafner et al., 2019) in which the PlaNet benchmark is ﬁrstly introduced.
To further evaluate the quality of the representations obtained by LCER, we also consider the statebased SAC (which called StateSAC) that directly learns from internal physical states provided by
DMControl Suite.
The performance of StateSAC can be seen as an upper bound since it utilizes the oracle representations of the environments.
All parameters of StateSAC are the same with LCER except that it uses states as input instead of pixels.
In Figure 1, we report the learning curves of LCER, PI-SAC (Lee et al., 2020b), SLAC
(Lee et al., 2020a), and PlaNet (Hafner et al., 2019).
LCER consistently achieves better or comparable performance than other methods.
In Table 1, We also report the performance of each algorithm at 100k environment steps.
LCER is the only algorithm that enables pixel-based SAC to outperform statebased SAC w.r.t the average returns, showing that the obtained representations can match the oracle descriptions( i.e. the physical states) of the environment.
In practice, pixel-based observations are often equipped with data augmentation to serve as a regularizer (Kostrikov et al., 2020b; Lee et al., 2020b; Laskin et al.; 2020).
In LCER, we also randomly shift the observations by [ − 4 , 4] pixels.
However, a proper data augmentation method is not always available for general forms of observations.
Therefore, We further test our method LCER in the situation where no data augmentation is available.
As shown in Figure 2, LCER can still improve the performance to a great extent without the help of data augmentation.
Note that some sampleefﬁcient RL algorithms such as DrQ (Kostrikov et al., 2020b) and RAD (Laskin et al., 2020) that rely on data augmentation can not be extended to such a situation naturally.
Our experiments on the PlaNet benchmark show that our proposed method LCER can obtain representations that help solve the tasks, and thus improve the sample efﬁciency.
5. R OBUSTNESS TO D ISTRACTIONS
In this section, we test LCER on the Distracting Control Suite (Stone et al., 2021), a challenging benchmark that is implemented by introducing several kinds of visual distractions to the DMControl
Suite, such as random changes to camera pose, object color, and background (see Figure 3).
These distractions widely exist in the real world and can reduce the performance of RL algorithms to a great extent (Stone et al., 2021).
We aim to show LCER’s ability to ﬁlter out those distractions and that it is not enough to just encode predictive information.
Note that some representation learning methods such as SLAC (Lee et al., 2020a) rely on a reconstruction loss, which will lead to signiﬁcant performance degradation as shown in (Zhang et al., 2020).
In Figure 3, we run experiments on three control tasks from the Distracting Control Suite: D.cheetahrun, D.walker-walk, D.ﬁnger-spin 5 .
In all tasks, we set the difﬁculty to be ‘easy’ and use a dynamically varying background (refer to (Stone et al., 2021) for more details).
We choose β = 0 .
4 for
D.cheetah-run and D.ﬁnger-spin, and β = 1 for D.walker-walk 6 .
As shown in Figure 3, proper representation learning methods are crucial in these tasks, because it is difﬁcult for the agent to ﬁlter out task-irrelevant information by itself.
LCER achieves much better performance than PI-SAC, which indicates that it can not lead to high-quality representations just by encoding predictive information, instead we have to attach more importance to the controllable elements.
We argue that this is because the distractions are often not controllable by the agent, and thus can be ﬁltered out by LCER.
Although these distractions do have a negative effect on LCER’s performance, LCER suffer less than other methods, showing LCER’s robustness to these distractions.
5. T HE E FFECT OF K AND β
In this section, we try to specify the effect of K and β , which are the most important hyperparameters in LCER.
We sweep over { 0.1,0.2,0.4,1 } to get the optimal β .
The learning curves for these β ’s are provided in Figure 4(b). [Figure] Figure β Figure 4(a) shows that larger K generally leads to better performance.
The effect of β is more complicated as shown in Figure 4(b). Larger β dose lead to better performance within a certain range, however when β goes beyond this range( i.e.β = 1 ), it may hinder the performance. In practice, we can gradually increase the value of β to search for optimal performance. To investigate the effect of K , we set β = 0 and report the performance for K = 0 , 1 , 2 , 5 in Figure 4(a).
Note that K = 0 is a special case in LCER that does not take into account any actions, and exactly recovers the work of CURL (Laskin et al.).
We add K = 0 to emphasize the importance of actions.
As shown in Figure 4(a), larger K leads to better performance, especially when K changes from 0 to 1.
We think this is because large K can provide more supervision signals thanks to our special design of Eq.(11).
To investigated the effect of β , we report the performance for β = 0 . 1 , 0 . 2 , 0 . 4 , 1 in Figure 4(b).
As shown in Figure 4(b), larger β dose lead to better performance within a certain range, however when β goes beyond this range ( e.g.β = 1 ), it may hinder the performance in some tasks.
The result in
Figure 4(b) indicates that the optimal β should be chosen according to the given environment.
By introducing a hyper-parameter β , LCER is able to strike a balance on the proportion of controllable elements that φ captures, allowing for searching for better performance.
5. T HE Q UALITY OF R EPRESENTATIONS
In this section, we try to measure the quality of φ by generalizing φ to unseen tasks.
Speciﬁcally, we train φ on a source task, froze φ , and then apply it to a target task within the same domain.
Since tasks within the same domain share the same environment dynamics, the performance on target tasks can measure how well φ can reﬂect the dynamics, and thus can be a metric of the representations’ quality.
In Figure 5, we compare φ (“LCER”) with another representation function ψ (“SAC transferred”) that is learned by the vanilla SAC.
We ﬁrst train φ or ψ on source tasks for 300k environment steps and then apply it to target tasks with frozen weights.
As shown in Figure 5, φ generally achieves better performance than ψ , which indicates that LCER can learn better representation than the vanilla SAC.
This is because φ can reﬂect the environment’s dynamics which are shared across the tasks from the same domain, while ψ only encodes those elements that are helpful to solve the source task.
As shown in Figure 5, φ can also accelerate the training process on target tasks, showing LCER’s potential to transfer learning.
C ONCLUSION
In this paper, we introduce the LCER, an efﬁcient framework to train representations that can reﬂect the action-related dynamics of the environment.
By focusing on the controllable elements of the environment, LCER can encode task-relevant information into representations, improving the performance of RL algorithms.
Our experiments on DMControl 100K benchmark show that LCER can enable the pixel-based SAC to outperform state-based SAC, demonstrating improvement over prior works.
We also evaluate LCER on the Distracting Control Suite, showing LCER’s robustness to distractions.
A I MPLEMENTATION
In this section, we explain the implementation in details.
Action Repeat Following (Lee et al., 2020b; Laskin et al.
; Lee et al., 2020a; Hafner et al., 2019), we use different action repeat to different environments.
The action repeat is a parameter that determine how frequently the agent is supposed to make a decision.
For a fair comparison, we adopt the action repeat from (Lee et al., 2020b;a).
Speciﬁcally, for cheetah-run, cartpole-swingup, reachereasy, and ball in cup-catch, we use action repeat of 4.
For walker-walk and ﬁnger-spin, we useand 1 respectively.
Note that the action repeat parameters can directly affect the performance.
Initial collecting step Following (Lee et al., 2020b;a), we pre-train our model using a small amount of random data.
Speciﬁcally, we ﬁrst interact with the environment using a random policy to collect initial experience, and then using the experience to pre-train our encoder φ .
Note that the amount of initial random data has been taken into account in all experiments.
For cheetah-run, walker-walk, reacher-easy, and ball in cup-catch, we collect random data for 10000 true environment steps and then train φ for 10000 iterations.
For ﬁnger-spin, 10000 true environment steps and 50000 iterations respectively.
For cartpole-swingup, 4000 true environment steps and 1000 iterations.
Network Architectures We use the encoder( i.e. φ ) that shares the same architecture with (Laskin et al.
; Lee et al., 2020b), which consists of 4 convolution layers(32 3 × 3 ﬁlters with stride 2 for the ﬁrst layer, and 32 3 × 3 with stride 1 for the rest 3 layers) and a fully-connected layer that ﬂattens the output of convolution layers and then maps it into a latent space with 50 dimensions.
For the h in
Eq.(8), we use MLPs with 2 128-d hidden layers, whose input and output are both of 50 dimensions.
For the actor and critic, we both use MLPs with two 512-d hidden layers.
We use tanh to squash the output of the actor to make sure it ﬁts the action bound of the environment
[Table]
Other Hyperparameters We list the other hyperparamters in Table 2.
Most of these paramters are the same with (Lee et al., 2020b).
[Table]
B E STIMATION OF C ONDITIONAL M UTUAL I NFORMATION
In this section, we describe how we optimize I JSD deﬁned in Section 4.2.
In the followings, the φ ( S t + k ) , A t : t + k − 1 , φ ( S t ) are denoted by X, Y, Z respectively
[Equation]
We ﬁrst borrow some theoretical results from (Nowozin et al., 2016):
[Equation]
Here, f JSD is the generator function of Jensen-Shannon divergence: f JSD ( u ) = − ( u +1) log 1+ u 2 + u log u
[Equation]
and f ∗ JSD is its corresponding convex conjugate function: f ∗ JSD ( t ) =
sup u { ut − f JSD ( u ) } = − log(2 − exp( t )) .
We choose V ( x, y, z ) = log(2)
− log(1 + exp(
− T k
( x, y, z )))
, then:
I
[Equation]
In Eq.(13), we suppose p ( y | z )
≈ p ( y ) ( i.e. A t : t + k − 1 is independent from φ ( S t ) ) for simplicity.
Such an assumption is appropriate in our case, because ( S t , A t : t + k − 1 , S t + k ) is sampled from a large replay buffer( = 10 5 ) whose action policy is quite non-stationary due to policy updates.
To estimate the RHS of Eq.(13), we ﬁrst randomly sample B transitions { s it , a it : t + k − 1 , s it + k } Bi
= from the replay buffer.
Assume ( j (1) , .., j ( B )) is a permutation of (1 , 2 , .., B ) obtained by random shufﬂing, then E x,y,z ∼ XY Z ( · ) can be approximated by 1 B (cid:80)
s it ,a it : t + k − 1 ,s it + k ( · ) and E x,z ∼ XZ,y ∼ Y ( · ) by 1 B (cid:80) s it ,a j ( i ) t : t + k − 1 ,s it + k ( · ) .
Therefore, the RHS of Eq.(13) can be estimated by:
[Equation]
Here, we omit the “ 2 log 2 ” in Eq.(13) for simplicity.