Adversarial Intrinsic Motivation for Reinforcement Learning

1 Introduction . Reinforcement Learning ( RL ) [ 74 ] deals with the problem of learning a policy to accomplish a given task in an optimal manner . This task is typically communicated to the agent by means of a reward function . If the reward function is sparse [ 4 ] ( e.g. , most transitions yield a reward of 0 ) , much random exploration might be needed before the agent experiences any signal relevant to learning [ 11 , 2 ] . Some of the different ways to speed up reinforcement learning by modifying or augmenting the reward function are shaped rewards [ 52 ] , redistributed rewards [ 2 ] , intrinsic motivations [ 8 , 69 , 71 , 72 , 54 , 57 ] , and learned rewards [ 81 , 54 ] . Unfortunately , the optimal policy under such modified rewards might sometimes be different than the optimal policy under the task reward [ 52 , 18 ] . The problem of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . learning a reward signal that speeds up learning by communicating what to do but does not interfere by specifying how to do it is thus a useful and complex one [ 82 ] . This work considers whether a task-dependent reward function learned based on the distribution mismatch between the agent ’ s state visitation distribution and a target task ( expressed as a distribution ) can guide the agent towards accomplishing this task . Adversarial methods to minimize distribution mismatch have been used with great success in generative modeling [ 29 ] and imitation learning [ 39 , 25 , 79 , 76 , 28 ] . In both these scenarios , the task is generally to minimize the mismatch with a target distribution induced by the data or expert demonstrations . Instead , we consider the task of goal-conditioned RL , where the ideal target distribution assigns full measure to a goal state . While the agent can never match this idealized target distribution perfectly unless starting at the goal , intuitively , minimizing the mismatch with this distribution should lead to trajectories that maximize the proportion of the time spent at the goal , thereby prioritizing transitions essential to doing so . The theory of optimal transport [ 78 ] gives us a way to measure the distance between two distributions ( called the Wasserstein distance ) even if they have disjoint support . Previous work [ 3 , 32 ] has shown how a neural network approximating a potential function may be used to estimate the Wasserstein-1 distance using its dual formulation , but assumes that the metric space this distance is calculated on is Euclidean . A Euclidean metric might not be the appropriate metric to use in more general RL tasks however , such as navigating in a maze or environments where the state features change sharply with transitions in the environment . This paper introduces a quasimetric tailored to Markov Decision Processes ( MDPs ) , the time-step metric , to measure the Wasserstein distance between the agent ’ s state visitation distribution and the idealized target distribution . While this time-step metric could be an informative reward on its own , estimating it is a problem as hard as policy evaluation [ 31 ] . Instead , we show that the dual objective , which maximizes difference in potentials while utilizing the structure of this quasimetric for the necessary regularization , can be optimized through sampled transitions . We use this dual objective to estimate the Wasserstein-1 distance and propose a reward function based on this estimated distance . An agent that maximizes returns under this reward minimizes this Wasserstein-1 distance . The competing objectives of maximizing the difference in potentials for estimating the Wasserstein distance and minimizing it through reinforcement learning on the subsequent reward function leads to our algorithm , Adversarial Intrinsic Motivation ( AIM ) . Our analysis shows that if the above Wasserstein-1 distance is computed using the time-step metric , then minimizing it leads to a policy that reaches the goal in the minimum expected number of steps . It also shows that if the environment dynamics are deterministic , then this policy is the optimal policy . In practice , minimizing the Wasserstein distance works well even when the environment dynamics are stochastic . Our experiments show that AIM learns a reward function that changes smoothly with transitions in the environment . We further conduct experiments on the family of goal-conditioned reinforcement learning problems [ 1 , 63 ] and show that AIM when used along with hindsight experience replay ( HER ) greatly accelerates learning of an effective goal-conditioned policy compared to learning with HER and the sparse task reward . Further , our experiments show that this acceleration is similar to the acceleration observed by using the actual distance to the goal as a dense reward . 2 Related Work . We highlight the related work based on the various aspects of learning that this work touches , namely intrinsic motivation , goal-conditioned reinforcement learning , and adversarial imitation learning . 2.1 Intrinsic Motivation . Intrinsic motivations [ 8 , 57 , 56 ] are rewards presented by an agent to itself in addition to the external task-specific reward . Researchers have pointed out that such intrinsic motivations are a characteristic of naturally intelligent and curious agents [ 30 , 5 , 6 ] . Intrinsic motivation has been proposed as a way to encourage RL agents to learn skills [ 10 , 9 , 68 , 62 ] that might be useful across a variety of tasks , or as a way to encourage exploration [ 11 , 67 , 7 , 24 , 23 ] . The techniques that encourage general exploration [ 16 , 58 ] have also been studied as inducing an artifical curiosity [ 66 , 65 ] . The optimal reward framework [ 69 , 72 ] and shaped rewards [ 52 ] ( if generated by the agent itself ) are a way to assist an RL agent in learning the optimal policy for a given task . Such an intrinsically motivated reward signal has previously been learned through various methods such as evolutionary techniques [ 54 , 64 ] , meta-gradient approaches [ 71 , 81 , 82 ] , and others . The Wasserstein distance has been used to present a valid reward for imitation learning [ 79 , 19 ] as well as program synthesis [ 27 ] . 2.2 Goal-Conditioned Reinforcement Learning . Goal-conditioned reinforcement learning [ 43 ] can be considered a form of multi-task reinforcement learning [ 17 ] where the agent is given the goal state it needs to reach at the beginning of every episode , and the reward function is sparse with a non-zero reward only on reaching the goal state . UVFA [ 63 ] , HER [ 1 ] , and others [ 80 , 20 ] consider this problem of reaching certain states in the environment . Relevant to our work , Venkattaramanujam et al . [ 77 ] learns a distance between states using a random walk that is then used to shape rewards and speed up learning , but requires goals to be visited before the distance estimate is useful . DisCo RL [ 51 ] extends the idea of goal-conditioned RL to distribution-conditioned RL . Contemporaneously , Eysenbach et al . [ 21 , 22 ] has proposed a method which considers goals and examples of success and tries to predict and maximize the likelihood of seeing those examples under the current policy and trajectory . For successful training , this approach needs the agent to actually experience the goals or successes . Their solution minimizes the Hellinger distance to the goal , a form of f -divergence . AIM instead uses the Wasserstein distance which is theoretically more informative when considering distributions that are disjoint , and does not require the assumption that the agent has already reached the goal through random exploration . Our experiments in fact verify the hypothesis that AIM induces a form of directed exploration in order to reach the goal . 2.3 Adversarial Imitation Learning and Minimizing Distribution Mismatch . Adversarial imitation learning [ 39 , 25 , 79 , 76 , 28 ] has been shown to be an effective method to learn agent policies that minimize distribution mismatch between an agent ’ s state-action visitation distribution and the state-action visitation distribution induced by an expert ’ s trajectories . In most cases this distribution that the expert induces is achievable by the agent and hence these techniques aim to match the expert distribution exactly . In the context of goal-conditioned reinforcement learning , GoalGAIL [ 20 ] uses adversarial imitation learning with a few expert demonstrations to accelerate the learning of a goal-conditioned policy . In this work , we focus on unrealizable target distributions that can not be completely matched by the agent , and indeed , are not induced by any trajectory distribution . FAIRL [ 28 ] is an adversarial imitation learning technique which minimizes the Forward KL divergence and has been shown experimentally to cover some hand-specified state distributions , given a smoothness regularization as used by WGAN [ 32 ] . f -IRL [ 53 ] learns a reward function where the optimal policy matches the expert distribution under the more general family of f -divergences . Further , techniques beyond imitation learning [ 46 , 37 ] have looked at matching a uniform distribution over states to guarantee efficient exploration . 3 Background . In this section we first set up the goal-conditioned reinforcement learning problem , and then give a brief overview of optimal transport . 3.1 Goal-Conditioned Reinforcement Learning . Consider a goal-conditioned MDP as the tuple 〈S , A , G , P , ρ0 , σ , γ〉 with discrete state space S , discrete action space A , a subset of states which is the goal set G ⊆ S , and transition dynamics P : S×A× G 7 −→ ∆ ( S ) ( ∆ ( · ) is a distribution over a set ) which might vary based on the goal ( see below ) . ρ0 : ∆ ( S ) is the starting state distribution , and σ : ∆ ( G ) is the distribution a goal is drawn from . γ ∈ [ 0 , 1 ) is the discount factor . We use discrete states and actions for ease of exposition , but our idea extends to continuous states and actions , as seen in the experiments . At the beginning of an episode , the starting state is drawn from ρ0 and the goal for that episode is drawn from σ . The reward function r : S × A × S × G 7 −→ R is deterministic , and r ( st , at , st+1|sg ) : = I [ st+1 = sg ] . That is , there is a positive reward when an agent reaches the goal ( st+1 = sg ) , and 0 everywhere else . Since the goal is given to the agent at the beginning of the episode , in goal-conditioned RL the agent knows what this task reward function is ( unlike the more general RL problem ) . The transition dynamics are goal-conditioned as well , with an automatic transition to an absorbing state s̄ on reaching the goal sg and then staying in that state with no rewards thereafter ( P ( s̄|sg , a , sg ) = 1 ∀ a ∈ A and P ( s̄|s̄ , a , sg ) = 1 ∀ a ∈ A ) . In short , the episode terminates on reaching the goal state . The agent takes actions in this environment based on a policy π ∈ Π : S × G 7 −→ ∆ ( A ) . The return Hg for an episode with goal sg is the discounted cumulative reward over that episode Hg =∑∞ t=0 γ tr ( st , at , st+1|sg ) , where s0 ∼ ρ0 , at ∼ π ( ·|st , sg ) , and st+1 ∼ P ( ·|st , at , sg ) . The agent aims to find the policy π∗ = arg maxπ∈Π Eg∈GEs0∼ρ0Eπ [ Hg ] that maximizes the expected returns in this goal-conditioned MDP . For a policy π , the agent ’ s goal-conditioned state distribution ρπ ( s|sg ) = Es0∼ρ0 [ ( 1− γ ) ∑∞ t=0 γ tP ( st = s|π , sg ) ] . Overloading the terminology a bit , we also define the goal-conditioned target distribution ρg ( s|sg ) = δ ( sg ) , a Dirac measure at the goal state sg . While learning using traditional RL paradigms is possible in goal-conditioned RL , there has also been previous work ( Section 2.2 ) on leveraging the structure of the problem across goals . Hindsight Experience Replay ( HER ) [ 1 ] attempts to speed up learning in this sparse reward setting by taking episodes of agent interactions , where they might not have reached the goal specified for that episode , and relabeling the transitions with the goals that were achieved during the episode . Off-policy learning algorithms are then used to learn from this relabeled experience .

Summary: This paper proposes Adversarial Intrinsic Motivation (AIM), a state distribution matching algorithm for goal-conditioned reinforcement learning. State distribution matching is an effective approach to imitation learning where the target distribution is yielded by the expert behavior. However, it is unclear how to apply it to goal-conditioned RL because the target distribution is unachievable - it puts all mass on the goal states and none on the trajectories leading to them. This paper proposes to use Wasserstein-1 distance as the distance measure between the agent's state visiting distribution and the goal distribution. To measure the distance between two states, the authors propose a quasimetric specific to MDPs - the time-step metric. The authors then prove that the policy that minimizes the Wasserstein-1 distance between the agent's state visiting distribution and the goal distribution under the time-step metric also minimizes the steps needed to achieve the goal. Empirical results show that AIM can provide dense and effective reward signals to the agent so that it learns better than using only the sparse reward of achieving the goal. When combined with Hindsight Experience Replay (HER), AIM outperforms the sparse goal-achieving rewards in a goal-conditioned RL setup and approaches the performance of a near-oracle hand-designed reward signal.