Wasserstein Autoencoders | Florent Delgrange

The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models

Wed, 29 May 2024 00:00:00 +0000

WBU

Fri, 10 May 2024 00:00:00 +0000

ICLR 2024 paper: https://openreview.net/forum?id=KrtGfTGaGe

The source code is available on GitHub.

Wasserstein Belief Updater (WBU) is an RNN free RL algorithm for POMDPs that learns a representation of the history via an approximation of the belief update in a reliable latent space model, providing theoretical guarantees for learning the optimal value.

This work concerns agents learning how to behave, i.e., their control policy, through reinforcement learning (RL). In real-world scenarios, the environment’s state is very often perceived either through noisy sensors, cameras, or more geneally imperfect observations (e.g., visual observation vs. exact coordinates on a map). In that case, the observation is non-Markovian and the environment is partially observable. This usually leads to complications compared to theoretical perfect-observation RL (i.e., with Markovian observation). For optimal decision making, the agent must in that case base its decision either on (a) the full observation-action history, or (b) the distribution over the possible real states of the environment in which the agent could be at each time step. The latter is called the belief of the agent and is a sufficient statistic to optimize the agent’s return.

The easiest method to tackle partial observability is to process the full history through an RNN to obain a compressed hidden state that can be fed to the policy of the learning agent. While appealing, RNNs don’t yield any guarantee that the representation learned is actually useful (a sufficient statistic) to optimize the agent’s return.

With WBU, we rather propose to learn a representation of the belief. Belief learning is difficult in RL because (1) the dynamics of the environments must be known to exactly compute the belief, and (2) it does not scale as it requires to integrate over the full state space (usually intractable). To tackle those challenges, WBU

learns a world model, through Wasserstein auto-encoded MDPs. This model comes with theoretical abstraction quality guarantees. It is learned through discrete latent spaces which eases the computation of the belief through the latent space.
minimizes the discrepancy between the theoretical belief update rule and the latent belief computed. This yields theoretical representation quality guarantees: close points in the representation space of the beliefs are guaranteed to yield close expected returns (Lipschitz continuity). This guarantees to support policy learning.

WAE-PCN: Wasserstein-autoencoded Pareto Conditioned Networks

Mon, 29 May 2023 00:00:00 +0000

WAE-MDPs

Wed, 01 Mar 2023 00:00:00 +0000

Official implementation of the Wasserstein Auto-encoded MDP (WAE-MDP) framework. WAE-MDPs enable the distilation of (any kind of) reinforcement learning policies into simpler controllers, paired with a discrete, tractable model of the environment (a latent space model). The two are provided with bisimulation guarantees, which allow formally verifying the behaviors of the agent operating under the simplified policy. The source code provided allows for replicating the experiments of the paper Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees.

The source code is available on GitHub.

Deep RL generally lacks guarantees but are very effective to solve complex tasks in intricate environments. On the other hand, model-checking techniques allow to verify agent’s policies, but need the access to an environment’s model. Since model-checking usually relies on an exhaustive exploration of the input model, the latter must in general be finite (discrete spaces) and tractable.

WAE-MDPs distill RL policies learned via any RL technique (in particular, through deep neural networks) to learn a discrete, latent representation of the intractable input space. The distillation yields both a simpler policy and a world model that are amenable to model-checking.

One of the main features of WAE-MDPs is that the quality of the abstraction can be formally checked through PAC bounds on bisimulation pseudometrics. In short, the world model is guaranteed to be bisimilarly close to the real environment; it replicates closely the agent-environment interaction behaviors through its latent space. In addition, WAE-MDPs are equipped with representation guarantees: the representation learned ensures that input states that are grouped to the same latent state are bisimilarly close (the agent behaves the same way from those states).

Beyond those PAC-verifiable guarantees, the bisimulation guarantees are also maintained during learning. This is in contrast to VAE-MDPs, which learn a surrogate, variational objective to avoid the intractable computation of the Wasserstein metric, in the fixpoint definition of bisimulation. Instead, WAE-MDPs soundly approximate this computation by learning a discriminator (“max”) that distinguishes between states that are likely to be produced in the learned world model, and those of the real environment. This yields a $2$-player game objective (min/max) where the goal of the minimizer is to learn both a representation and a world model that makes the latent states generated by the encoder and the world model indistinguishable.

By doing so, WAE-MDPs avoid posterior collapse issues, which drastically speeds up learning, and yield distilled policies and world models of better quality.

Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees

Sun, 01 Jan 2023 00:00:00 +0000