Cell-Free Latent Go-Explore

This paper uses the standard formalism of a discounted Markov Decision Process (MDP) defined as the tuple $(S,A,P,R,\gamma ,{\rho }_0)$ where $S$ is the set of states, $A$ is the set of actions, $P:S\times A\to S$ is the (unknown) transition function providing the probability distribution of the next state given a current state and action, $R:S\times A\times S\to R$ is the reward function, $\gamma$ is the discount factor and ${\rho }_0$ is the initial distribution of states. A policy, denoted $\pi (\cdot |\cdot ):S\times A\to R^{+}$ is the probability distribution such that $\pi \left(s|a\right)$ is the probability of selecting action $a\in A$ in state $s\in S$. We denote the previously defined values with discrete time $t$ such that $s_t$, $a_t$ and $r_t$ denote respectively the state, action, and reward at timestep $t$. The aim is to learn a policy $\pi$ that maximises the long-term expected reward $E\left[{\sum }_{t=0}^{+\infty }{\gamma }^t{r}_t\right]$.

We note that every MDP can be augmented into a goal-conditioned MDP with a goal space $G$ and an initial goal distribution ${\rho }_g$. At each timestep, the observation is augmented with a goal and the reward function depends on this goal. A goal-conditioned policy kaelbling1993learning, denoted $\pi (\cdot |\cdot ,\cdot )$ also depends on the goal.

In random exploration, the agent does not follow a specific exploration approach. The agent’s actions can be purely random, or in the continuous case, augmented by exploration noise parametrized by the current state haarnoja2018soft or not lillicrap2016continuous.

Intrinsic rewards-based methods are inspired by intrinsic motivation from cognitive science oudeyer2009intrinsic. They consist in adding an additional reward (called intrinsic) to the reward signal from the environment (called extrinsic). This intrinsic reward is intended to foster exploration. It can be based on the state visitation count bellemare2016unifying; machado2020count or on the prediction error of a model learned from the collected data houthooft2016vime; pathak2017curiosity; achiam2017surprise; pathak2019self; burda2019exploration; tao2020novelty.

Methods that modify the reward of the environment provide no mechanism for distilling the knowledge gained from visiting various states. Agents visit new states, and they quickly forget them when other states become newer. In contrast, recent work suggests the use of a goal-conditioned autotelic agent specifically trained for the exploration task allows to use this knowledge for the realization of new user-specified goals levine2021understanding; colas2022autotelic. The reward signal is ignored during the exploration phase. After the exploration phase, the data collected by the agent are used to learn one or more subsequent tasks jin2020reward. In goal-based methods, the agent is conditioned by a goal that is used to guide exploration towards unknown areas. These methods rely on a goal generator to create goals for the agent. We divide goal-based methods into two categories: exploratory goal methods and goals to explore from methods.

Exploratory goal methods follow the intuition that the agent discovers new areas of the observation space by pursuing goals that have been little or not achieved before. The challenge of these methods is to choose the goal to be neither too easy nor too hard. The literature contains several ways to approach this trade-off. Some methods sample goals that either maximize Learning Progress colas2019curious; portelas2020teacher or value disagreement zhang2020automatic. Other methods sample goals from the least visited areas using a parametric density model on the visited states pong2020skew. It is also possible to imagine goals that have never been reached using a language model colas2020language, a generative model racaniere2020automated or a GAN florensa2018automatic.

In goals to explore from methods the agent sample a goal from previously visited states. It returns to it, either by teleportation ecoffet2019go; matheron2020efficient, or with goal-conditioned policy ecoffet2021first. The challenge of these methods is to choose a goal that is of high exploratory interest. Similarly, some methods estimate the density of the encountered states, either with parametric methods pitis2020maximum or with non-parametric methods ecoffet2021first; matheron2020efficient, to target the low-density areas.

In summary, methods based on a goal reaching policy should facilitate scalable RL. The stunning results of Go-Explore illustrate this point but remain circumscribed to few environments and require a lot of domain knowledge to work. By bridging with concepts already used in the intrinsic reward literature, we show a way to make this approach more general and simple.

pathak2017curiosity proposed an intrinsic reward calculated based on the agent’s prediction error of the consequence of its own actions. The representation is learned using two submodules. The first encodes the observation into a latent representation $\varphi \left(s_t\right)$. The second takes as input $\varphi \left(s_t\right)$ and $\varphi (s_t+1)$ and outputs the prediction of the action taken by the agent at time step $t$. The parameters $\theta$ of the inverse model $P_{\theta }^{inv}$ are optimized by minimizing the loss function:

The inverse dynamics representation learning allows getting a latent representation of the states containing only the aspects of the state on which the agent can have an influence.

In achiam2017surprise, the intrinsic reward is calculated based on the prediction error of a model approximating the transition probability function of the MDP. Two submodules are used. The first one encodes the observation to a latent representation $\varphi \left(s_t\right)$. The second takes as input $\varphi \left(s_t\right)$ and $a_t$ and outputs the prediction of the next state ${\hat{s}}_{t+1}$. The model parameters $\theta$ are optimized by minimizing the loss function:

We train an agent to navigate in a continuous 2D maze. The corresponding configuration is shown in Figure 5. The agent is initially at the center of the maze. At each timestep, the agent receives the current coordinates as an observation and chooses an action that controls its location change. If the agent collides with a wall, it returns to its previous position. The reachable space is a square of $12\times 12$ and the agent’s action is limited to $[-1,1]$ horizontally and vertically. The agent can interact 100 times with the environment, after which the episode ends and the agent returns to its initial position in the center of the maze.

Robotic environments are interesting and challenging application cases of RL, especially since the reward is often sparse. We simulate a Franka robot under the PyBullet physics engine using panda-gym gallouedec2021pandagym. Figure 3 depicts a rendering of the setup. The robot can move and interact with an object. The agent has access to the position of the end-effector and the position of the object, as well as to the opening of the gripper. The agent interacts 50 times with the environment and then the object and the robot arm are reset to their initial position.

Most RL methods from the literature do not follow any specific exploration strategy. In a reward-free context, the performance of the latter is often equivalent to a random walk. We will take as a baseline a random agent, whose actions are uniformly sampled over the action space at each timestep, Soft Actor-Critic (SAC; haarnoja2018soft) and Deep Deterministic Policy Gradient (DDPG; lillicrap2016continuous).

In this paper, we take as reference two widely used intrinsic reward-based methods combined with either SAC or DDPG. These methods stand out from the others because, despite their simplicity, they have demonstrated good performance on a wide variety of tasks.

• Intrinsic Curiosity Module (ICM; pathak2017curiosity): the intrinsic reward is computed as the mean square error between the true latent representation and the one predicted by a learned dynamic model given the action taken. The encoder is trained jointly with an inverse dynamics model.

• Surprise achiam2017surprise: the intrinsic reward is the approximation of the KL divergence between the actual transition probabilities and a learned transition model.

We use the following algorithms as a baseline.

• Go-explore ecoffet2021first: the agent discretizes the state space into cells, selects the least visited ones, returns to them using a goal-conditioned policy, and further explores from there. We use a constant cell generation strategy during learning, like the one proposed in the original paper for robotic tasks. Go-Explore is the closest baseline to our algorithm. Appendix D details the differences between LGE and Go-Explore.

• Diversity Is All You Need (DIAYN; eysenbach2019diversity): the agent is conditioned by a skill and a discriminator predicts the skill pursued by the agent. The more the discriminator predicts with certainty the skill pursued, the bigger the reward. Conjointly, the discriminator is trained to maximize the distinguishability of skills.

• Skew-Fit pong2020skew: the agent’s goal sampling is skewed to maximize the entropy of a density model learned on the achieved states.

For the goal-directed methods, we use Hindsight Experience Replay (HER; andrychowicz2017hindsight) relabeling which has shown to significantly increase learning.

To nullify the variation in results due to different implementations, we implement all algorithms in the same framework: Stable-Baselines3. raffin2021stable.

The set of intrinsic reward-based methods and goal-directed methods are underpinned by the same off-policy algorithm. The hyperparameters for this algorithm are identical. For the Maze environment, we use SAC, while for the robotic environment, we use DDPG as it gives better results for all methods. To negate the influence of a bad choice of hyperparameter on the results, the method-specific hyperparameters are optimized. Appendix A details the optimization process for all methods and the resulting hyperparameters.

The approach we propose is based on the assumption that the agent is always initialized in the same state. This assumption guarantees that at the beginning of each episode, all the states previously reached are reachable and that the subgoal trajectory starts with the initial state of the agent. However, in some environments, especially in procedurally generated environments, this assumption is not fulfilled kuttler2020nethack. In this situation, trying to follow the subgoal trajectory may be counterproductive in reaching the final goal. It is also possible that the final goal is not even reachable. However, note that even the pursuit of an unreachable goal can foster exploration. Future work should investigate how to tailor the algorithm to overcome this problem. We believe that an approach inspired by generative networks such as racaniere2020automated may be appropriate.

In LGE, an agent is considered to have reached a goal (whether final or intermediate) when the latent distance between its state and the goal is below a specified threshold. This is a naive way of defining a goal achievement function colas2022autotelic that depends crucially on the latent representation. We believe that the results could be improved by envisioning a more informative and suitable goal achievement function for our method.

The main contribution of our work consists in the generalization of the Go-Explore approach by using a latent representation. LGE can be applied without further adaptation to high-dimensional environments. Our approach is particularly useful in environments where the observation is an image, and for which learning directly from the raw pixels is not desirable. Thus, representation learning is the keystone of the method. We provide a proof of concept for a forward and inverse model. We believe that these results can be greatly improved by choosing more finely the representation learning method for each environment by taking advantage of the many works dealing with this subject lesort2018state. Future work should confirm this intuition.