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 a method for goal-conditioned RL. The main idea is to learn a function that looks like a value-function and then use that value function as a reward function for RL. This method is motivated as minimizing a certain Wasserstein distance. Empirically, the proposed method outperforms goal-conditioned RL applied to sparse reward functions, and performs roughly on par with goal-conditioned RL with dense rewards. In effect, it seems like this learned value function can serve the role of the dense reward function.