V ERSATILE O FFLINE I MITATION L EARNING VIA S TATE O CCUPANCY M ATCHING
A BSTRACT
We propose State Matching Offline DIstribution Correction Estimation (SMODICE), a novel and versatile algorithm for offline imitation learning (IL) via state-occupancy matching.
Without requiring access to expert actions, SMODICE can be effectively applied to three offline IL settings: (i) imitation from observations (IfO), (ii) IfO with dynamics or morphologically mismatched expert, and (iii) example-based reinforcement learning, which we show can be formulated as a state-occupancy matching problem.
We show that the SMODICE objective admits a simple optimization procedure through an application of Fenchel duality, reducing a nested optimization problem to a sequence of stable supervised learning problems.
We extensively evaluate SMODICE on both gridworld environments as well as on high-dimensional offline benchmarks.
Our results demonstrate that
SMODICE is effective for all three problem settings and significantly outperforms prior state-of-art.
I NTRODUCTION
The offline reinforcement learning (RL) framework Lange et al.
; Levine et al.
(2020) aims to use pre-collected reusable offline data—without further interaction with the environment—for sample-efficient, scalable and practical data-driven decision-making.
However, this assumes that the offline dataset comes with reward labels, which may not always be possible.
To address this, offline imitation learning (IL) Zolna et al.
; Chang et al.
; Kim et al.
(2022) has recently been proposed as an alternative where the learning algorithm is provided with a small set of expert demonstrations and a separate set of offline data of unknown quality.
The goal is to learn a policy that mimics the provided expert data while avoiding test-time distribution shift Ross et al.
(2011) by using the offline dataset.
Expert demonstrations are often much more expensive to acquire than offline data; thus, offline IL benefits significantly from minimizing assumptions about the expert data.
In this work, we aim to remove two assumptions about the expert data in current offline IL algorithms: (i) expert action labels must be provided for the demonstrations, and (ii) the expert demonstrations are performed with identical dynamics (same embodiment, actions, and transitions) as the imitator agent.
These requirements preclude applications to important practical problem settings, including (i) imitation from observations, (ii) imitation with mismatched expert that obeys different dynamics or embodiment (e.g., learning from human videos), and (iii) learning only from examples of successful outcomes rather than full expert trajectories Eysenbach et al.
For these reasons, many algorithms for online IL have already sought to remove these assumptions Torabi et al.
Liu et al.
; Radosavovic et al.
Eysenbach et al.
but extending them to offline IL remains an open problem.
We propose S tate M atching O ffline DI stribution C orrection E stimation (SMODICE), a general offline IL framework that can be applied to all three problem settings described above.
At a high level, SMODICE is based on a state-occupancy matching view of IL; in particular, it optimizes a tractable offline upper bound of the KL-divergence of the state-occupancy d between the imitator π and the expert E :
[Equation]
This state-occupancy matching objective allows SMODICE to infer the correct actions from the offline data in order to match the state-occupancy of the provided expert demonstrations.
This naturally enables imitation when expert actions are unavailable, and even when the expert’s embodiment or dynamics are different, as long as there is a shared task-relevant state.
Finally, we show that example-based RL Eysenbach et al.
(2021), where only examples of successful states are provided as supervision, can be formulated as a state-occupancy matching problem between the imitator and a “teleporting” expert that is able to reach success states in one step.
Hence, SMODICE can also be used as an offline example-based RL 1 method without any modification.
[Figure]
Naively optimizing the offline upper bound on (1) would result in an actor-critic style IL algorithm akin to prior work Ho & Ermon (2016);
Kostrikov et al.
(2018; 2020); however, these algorithms suffer from training instability in the offline regime Kumar et al.
Lee et al.
; Kim et al.
(2022) due to the entangled nature of actor and critic learning, leading to erroneous value bootstrapping Levine et al.
SMODICE bypasses this issue and achieves “actor-free” training by directly estimating the importance weight ratio of the occupancy measures between the optimal policy and the empirical behavior policy of the offline data, leveraging the stationary distribution correction estimation Nachum et al. (2019a)
; Nachum & Dai (2020) (DICE) paradigm.
Specifically, by formulating the policy optimization problem via its dual (i.e., optimizing over the space of valid state-action occupancy distributions) and applying Fenchel duality, SMODICE obtains an unconstrained convex optimization problem over a value function arising from Lagrangian duality, which admits closed-form solutions in the tabular case and can be easily optimized using stochastic gradient descent (SGD) in the deep RL setting.
Then, SMODICE projects the optimal value function onto the dual space to extract the optimal importance weights, and learns the optimal policy via weighted Behavior Cloning.
Note that SMODICE does not learn a policy until the value function has converged.
Through extensive experiments, we show that SMODICE is effective for all three problem settings we consider and outperforms all state-of-art methods in each respective setting.
A benefit of the improved optimization is that SMODICE is substantially more stable compared to prior methods—we obtain all SMODICE results using a single set of hyperparameters, modulo a choice of f -divergence which can be tuned offline .
In contrast, prior methods suffer from much greater performance fluctuation across tasks and settings.
Altogether, our proposed method SMODICE can serve as a versatile offline
IL algorithm that is suitable for a wide range of assumptions on expert data.
Pedagogical examples.
To illustrate SMODICE’s versatility, we have applied it to two gridworld tasks, testing offline IL from mismatched experts and examples, respectively.
Figure 2(a) shows an expert agent that can move diagonally in any direction, whereas the imitator can only move horizontally or vertically.
In Figure 2(b), only a success state (the star) is provided as supervision.
An offline dataset collected by a random agent is given to SMODICE for training in both cases.
As shown, SMODICE recovers an optimal policy (i.e. minimum state-occupancy divergence to that of the expert) in both cases.
See Appendix D.1 for details.
P RELIMINARIES
Markov decision processes.
We consider a time-discounted Markov decision process (MDP) Puterman (2014)
M = ( S, A, R, T, µ 0 , γ ) with state space S , action space A , deterministic rewards
R ( s, a ) , stochastic transitions s ′ ∼ T ( s, a ) , initial state distribution µ 0 ( s ) , and discount factor γ ∈ (0 , 1] .
A policy π : S → ∆( A ) determines the action distribution conditioned on the state.
The state-action occupancies (also known as stationary distribution)
d π ( s, a ) : S × A → [0 , 1] of π is
[Equation]
which captures the relative frequency of state-action visitations for a policy π
The state occupancies then marginalize over actions:
d π ( s ) =
(cid:80) a d π ( s, a ) .
The state-action occupancies satisfy the single-step transpose Bellman equation:
[Equation]
where T π⋆ is the adjoint policy transition operator,
[Equation]
Divergences and Fenchel conjugates.
Due to space constraint, their definitions are included in
Appendix A.
Offline imitation learning.
Many imitation learning approaches rely on minimizing the f -divergence between the state-action occupancies of the imitator and the expert Ho & Ermon (2016);
Ke et al.
(2020); Ghasemipour et al.
[Equation]
In imitation learning, we do not have d E ; instead, we are provided with expert demonstrations
D E := { ( s ( i ) , a ( i ) ) } Ni =1 .
In offline imitation learning, the agent further cannot interact with the MDP M ; instead, they are given a static dataset of logged transitions D O := { τ i } Mi =1 , where each trajectory τ ( i ) =
( s ( i ) 0 , a ( i ) 0 , s ( i ) 1 , ... ) with s ( i ) 0 ∼ µ 0 ; we denote the empirical state-action occupancies of D O as d O ( s, a ) .
T HE SMODICE A LGORITHM
SMODICE aims to minimize a state-occupancy matching objective based on the KL-divergence:
[Equation]
[Equation]
Later, we show that we can optimize any f -divergence that upper bounds KL-divergence, as the conjugate of the KL-divergence may be numerically unstable.
Minimizing (6) requires on-policy samples from π , as the expectation is over d π .
To derive an offline objective, we first derive an upper bound involving the offline dataset.
First, we assume expert coverage of the offline data:
Assumption 1.
d O ( s ) > 0 whenever d E ( s ) >
This assumption ensures that the offline dataset has coverage over the expert state-marginal, and is necessary for imitation learning to succeed.
Whereas prior offline RL approaches Kumar et al.
Ma et al.
(2021) assume full coverage of the state-action space, our assumption 2 is considerably weaker since it only requires expert coverage.
Given this assumption, we proceed to derive the offline upper bound on state-occupancy matching:
Theorem 1.
Given Assumption 1, we have
[Equation]
Furthermore, for any f -divergence such that D f ≥ D KL ,
[Equation]
Proofs are in Appendix A. Intuitively, this says that offline state-occupancy matching can be achieved by matching states in the offline data that resemble expert states (the first term), while remaining in the support of the offline state-action distribution (the second term).
Replacing KL-divergence with other divergences can be useful since the conjugate of KL divergence involves a log-sum-exp, which has been found to be numerically unstable in many RL tasks Zhu et al.
; Lee et al.
Rudner et al.
Note that (8) requires samples from d π , so it still cannot be easily optimized without online interaction.
To address this, we first rewrite it as an optimization problem over the space of valid state-action occupancies Puterman [Equation] [Equation]
[Equation]
where T ⋆ d ( s ) = (cid:80) ¯ s, ¯ a T ( s | ¯ s, ¯ a ) d (¯ s, ¯ a ) ; here, (10) ensures that d is the occupancy distribution for some policy.
We assume that (9) is strictly feasible .
Assumption 2.
There exists at least one d ( s, a ) such that constraints (10) are satisfied and ∀ s ∈
S , d ( s ) > 0 .
This assumption is mild and can be satisfied in practice for any MDP for which every state is reachable from the initial state distribution.
Next, we can form the dual of (9):
[Equation]
where V ( s ) are the Lagrangian multipliers.
Now, because T ⋆ is the adjoint of T [Equation] we have the following:
[Equation]
Using this equation [Equation] we can write (11) as
[Equation]
We note that the original problem (9) is convex Lee et al.
By Assumption 2, it is strictly feasible, so by strong duality, we can change the order of optimization in (13):
[Equation]
[Equation]
where R ( s ) = log d E ( s )
d O ( s ) .
We briefly discuss how to compute R ( s ) .
In the tabular case, R ( s ) can be computed using empirical estimates of d E ( s ) and d O ( s ) .
In the continuous case, we can train a discriminator c : S → (0 , 1) :
[Equation]
The optimal discriminator is c ⋆ ( s ) =
[Equation]
et al
[Equation] , so we can use R ( s ) = − log
[Equation]
Finally, using the Fenchel conjugate, (14) can be reduced to a single unconstrained optimization problem over V : S → R ≥ 0 that depends on samples from only d O and not d ; we also obtain the importance weight of the state-occupancy of the optimal policy with respect to the offline data.
Theorem 2.
The optimization problem (14) is equivalent to
[Equation]
Furthermore, given the optimal solution V ∗ , the optimal state-occupancy importance weights are
[Equation]
This result can be viewed as using Fenchel duality to generalize prior DICE-based offline approaches Lee et al.
; Kim et al.
In particular, the inner maximization problem in (14) is precisely the Fenchel conjugate of D f ( d ( s, a )
∥
d O ( s, a )) at R ( s ) + γ T V ( s, a ) − V ( s ) (compare (14) to (23)).
Similarly, (17) can be derived from leveraging the relationship between the optimal solutions of a pair of Fenchel primal-dual problems.
This generality allows us to choose problemspecific f -divergences that improve stability during optimization.
In Appendix C, we specialize the
SMODICE objective for the KL- and χ 2 -divergences, which we use in our experiments.
Finally, using the optimal importance weights, we can extract the optimal policy π using weighted
Behavior Cloning:
[Equation] [Equation]
where ξ ∗ ( s, a )
=
d ∗ ( s,a ) d
O ( s,a ) .
Here, V ( s ) can be viewed as the value function—it is trained by minimizing a convex function of the Bellman residuals and the values of the initial states.
Then, it can be used to inform policy learning.
Putting everything together, SMODICE can achieve stable policy learning through a sequence of three disjoint supervised learning problems, summarized in Figure 1.
The full pseudo-code is in
Algorithm 2 in Appendix 2.
SMODICE for Tabular MDPs.
An appealing property of SMODICE is that it admits closed-form analytic solution in the tabular case.
The derivation is in Appendix D.
O FFLINE I MITATION L EARNING FROM E XAMPLES
Next, we describe how SMODICE can be applied to offline imitation learning from examples.
Starting from the original problem objective from Eysenbach et al.
(2021), we derive a state-occupancy matching objective, enabling us to apply SMODICE without any modification.
Problem setting.
We assume given success examples S ∗ =
{ s ∗ ∼ p
U ( s t | e t = 1) } , where e ∈ { 0 , 1 } indicates whether the current state is a success outcome, and
offline data D = { ( s, a, s ′ ) } .
Here, U is the state distribution of the “user” providing success examples.
Then, Eysenbach et al.
(2021) proposes the example-based RL objective
[Equation]
That is, we want a policy that maximizes the probability of reaching success states in the future.
To tackle this problem in the offline setting, our strategy is to convert (19) into an optimization problem over the state-occupancy space.
Derivation.
We first transform the problem into state-occupancy space—i.e.,
[Equation]
which is valid given that the original objective can be thought of as a regular RL problem with reward function r ( s )
=
p ( e | s )
Eysenbach et al.
Given this formulation, we can derive a tractable lower bound to (20) through Jensen’s inequality and
Bayes’ rule:
[Equation]
We can optimize the original objective by maximizing this lower bound.
Doing so is equivalent to solving min d ≥ 0 D KL ( d ( s ) ∥ p U
( s | e )) , which is exactly in the form of the state-occupancy matching objective (6) that SMODICE optimizes.
Furthermore, this objective admits an intuitive explanation from a purely imitation learning lens.
We can think of p U ( s | e ) as the state-occupancy distribution of an expert agent who can “teleport” to any success state in one time-step.
Therefore, we have shown that example-based RL can be understood as a state-occupancy minimization problem between a
MDP-dynamics abiding imitator and a teleporting expert agent.
Consequently, SMODICE can be used in the offline setting without any algorithmic modification.
E XPERIMENTS
We experimentally demonstrate that SMODICE is effective for offline IL from observations, mismatched experts, and examples.
We give additional experimental details in Appendices G, H, and
I.
5. O FFLINE I MITATION L EARNING FROM O BSERVATIONS
[Figure]
Datasets.
We utilize the D4RL Fu et al.
offline
RL dataset.
The dataset compositions for all tasks are listed in Table 3 in Appendix G.
We consider the following standard Mujoco environments: Hopper,
Walker2d, HalfCheetah, and Ant.
For each, we take a single expert trajectory from the respective “expertv2” dataset as the expert dataset and omit the actions.
For the offline dataset, following Kim et al.
(2022), we use a mixture of small number of expert trajectories ( ≤ 200 trajectories) and a large number of low-quality trajectories from the “random-v2” dataset (we use the full random dataset, consisting of around 1 million transitions).
This dataset composition is particularly challenging as the learning algorithm must be able to successfully distinguish expert from low-quality data in the offline dataset.
We also include two more challenging environments from D4RL: AntMaze and Franka Kitchen.
Dur to space constraint, we describe their dataset compositions in Appendix G.
Method and baselines.
We use SMODICE with χ 2 -divergence for all tasks (in other problem settings as well) except Hopper, Walker, and Halfcheetah, where we find SMODICE with KL-divergence to perform better; in Appendix E.2, we explain how to choose the appropriate f -divergence offline by monitoring SMODICE’s policy loss.
For comparisons, we consider both IfO and regular offline IL methods, which make use of expert actions.
For the former, we compare against (i) SAIL-TD3-BC , which combines a state-of-art state-matching based online IL algorithm (SAIL)
Liu et al.
with a state-of-art offline RL algorithm (TD3-BC) Fujimoto & Gu (2021), 3 (ii) Offline Reinforced
Imitation Learning (ORIL)
Zolna et al.
(2020), which adapts GAIL Ho & Ermon (2016) to the offline setting by using an offline RL algorithm for policy optimization
; we implement ORIL using the same state-based discriminator as in SMODICE, and TD3-BC as the offline RL algorithm.
For the latter, we consider the state-of-art DEMODICE Kim et al.
(2022) as well as Behavior Cloning (BC) .
We train all algorithms for 1 million gradient steps and keep track of the normalized score (i.e., 100 is expert performance, 0 is random-action performance) during training; the normalized score is averaged over 10 independent rollouts.
All methods are evaluated over 3 seeds, and one standard-deviation confidence intervals are shaded.
Results.
As shown in Figure 3, only SMODICE achieves stable and good performance in all six tasks.
It achieves (near) expert performance in all the Mujoco environments, performing on-par with DEMODICE and doing so without the privileged information of expert actions.
SMODICE’s advantage over DEMODICE is more apparent in AntMaze and Kitchen.
In the former, SMODICE outperforms BC, while DEMODICE cannot; in the latter, DEMODICE quickly collapses due to its use of KL-divergence, which may be numerically unstable in high-dimensional environments.
BC is a strong baseline for tasks where the offline dataset contains (near) expert data (i.e., AntMaze and
Kitchen); however, as the dataset becomes more diverse, BC’s performance drops significantly.
SAIL-TD3-BC and ORIL
both fail to learn in some environments and otherwise converge to a worse policy than SMODICE.
The only exception is AntMaze; however, in Appendix G.2, we show that both methods collapse with a more diverse version of the AntMaze offline dataset, indicating that unlike SMODICE, these methods are highly sensitive to the composition of the offline dataset, and work best with task-aligned offline data.
The sub-par performances of SAIL and ORIL highlight the challenges of adapting online IL methods to the offline setting; we hypothesize that it is not sufficient to simply equip the original methods (i.e., SAIL and GAIL) with a strong base offline RL algorithm.
Together, these results demonstrate that SMODICE is stable, scalable, and robust, and significantly outperforms prior methods.
Finally, in Appendix G.2, we ablate SMODICE by zeroing out its discriminator-based reward to validate that SMODICE’s empirical performance comes from its ability to discriminate expert data in the offline dataset.
5. O FFLINE IL FROM M ISMATCHED E XPERTS
Datasets and baselines.
Due to space constraint, they are described in Appendix H.
Results.
The training curves are shown in Figure 5; we illustrate the original maximum performance attained by each method (i.e., using the original expert trajectory, Section 5.1) using dashed lines as points of reference.
As can be seen, SMODICE is significantly more robust to mismatched experts than either SAIL-TD3-BC or ORIL.
On AntMaze, the task where SAIL-TD3-BC and ORIL originally outperform SMODICE, learning from a PointMass expert significantly deteriorates their performances, and the learned policies are noticably worse than that of SMODICE, which has the smallest performance drop.
The other two tasks exhibit similar trends; SMODICE is able to learn an expert level policy for the original Ant embodiment using a disabled Ant expert, and is the only method that shows any progress on the hardest HalfCheetah-Short task.
Despite using the same discriminator for reward supervision, SMODICE is substantially more robust than ORIL, likely due to the occupancy-constraint D f
( d ( s, a ) ∥
d O ( s, a )) term in its objective (8), which ensures that the learned policy is supported by the offline data as it attempts to match the expert states.
In Appendix
H.2, we provide additional quantitative analysis of Figure 5.
5. O FFLINE I MITATION L EARNING FROM E XAMPLES
Tasks.
We use the AntMaze and Kitchen environments and create example-based task variants.
For
AntMaze, we replace the full demonstration with a small set of success states (i.e., Ant in the goal region) extracted from the offline data.
For Kitchen, we consider two subtasks in the environment:
Kettle and Microwave.
and define task success to be only whether the specified object is correctly placed (instead of all objects as in the original task); the success states are extracted from the offline data accordingly.
Examples of the success states are illustrated in Figure 13 in Appendix I. Note that the kitchen dataset contains many trajectories where the kettle is moved first.
Thus, the kettle task is easy even for Behavior Cloning (BC), since cloning the offline data can lead to success.
This is not the case for the microwave task, making it much more difficult to solve using only success examples.
In addition, we introduce the PointMass-4Direction environment.
Here, a 2D PointMass agent is tasked with navigating to the middle point of a specified edge of the square that encloses the agent (see Figure 13(a)).
The offline dataset is generated using a waypoint navigator controlling the agent to each of the four possible goals and contains equally many trajectories for each goal.
At training and evaluation time, we set the left edge to be the desired edge and collect success states from the offline data accordingly.
This task is low-dimensional but consists of multi-task offline data, making it challenging for algorithms such as BC that do not solve the example-based RL objective.
Approaches.
We make no modification to SMODICE; the only difference is that the discriminator is trained using success states instead of full expert state trajectories.
Our main comparison is RCETD3-BC , which combines RCE Eysenbach et al.
(2021), the state-of-art online example-based RL method, and TD3-BC.
We also compare against ORIL Zolna et al.
(2020), using the same architecture as in Section 5.1.
Finally, we also include BC.
Results.
As shown in Figure 6, SMODICE is the best performing method on all four tasks and is the only one that can solve the Microwave task.
RCE-TD3-BC is able to solve the first three tasks, but achieves worse solutions and exhibits substantial performance fluctuation during training; we posit that the optimization for RCE, which requires alternate updates to a recursive classifier and a policy, is substantially more difficult than that of SMODICE.
ORIL is unstable and fails to make progress in most tasks.
Interestingly, as in the mismatched expert setting, on AntMaze, ORIL’s performance is far below that of SMODICE, despite attaining better results originally (Figure 3).
This comparison demonstrates the versatility of SMODICE afforded by its state-occupancy matching objective; in contrast, ORIL treats offline IL from examples as an offline RL task with discriminator-based reward and cannot solve the task.
See Appendix I for additional analysis.
C ONCLUSION
We have proposed SMODICE, a simple, stable, and versatile algorithm for offline imitation learning from observations, mismatched experts, and examples.
Leveraging Fenchel duality, SMODICE optimizes a state-occupancy matching objective that enjoys closed-form tabular solution and stable optimization with deep neural networks.
Through extensive experiments, we have shown that
SMODICE significantly outperforms prior state-of-art methods in all three settings.
A P ROOFS A. P RELIMINARIES
Definition 1 ( f -divergence) .
Given a continuous, convex function f and two probability distributions p, q ∈ ∆( X ) over a domain X , the f -divergence of p at q is
[Equation]
A common f -divergence in machine learning is the KL-divergence, which corresponds to f ( x ) = x log x .
Now, we introduce Fenchel conjugate for f -divergences.
Definition 2 (Fenchel conjugate) .
Given a vector space Ω with inner-product ⟨· , ·⟩ , the Fenchel conjugate f ⋆ : Ω ⋆
→ R of a function f : Ω → R is
[Equation]
For an f -divergence, under mild realizability assumptions Dai et al.
(2016) on f , the Fenchel conjugate of D f ( p ∥ q ) at y :
X → R is
[Equation]
A. T ECHNICAL L EMMAS
Lemma 1.
We have
[Equation]
Proof.
We first state and prove a related lemma, which first appeared in Yang et al.
Lemma 2.
[Equation]
Proof.
D KL
[Equation]
Using this result [Equation] we can show the desired upper bound:
[Equation]
A. P ROOF OF T HEOREM
Proof.
[Equation]
we assume that d O ( s ) > 0 whenever d E ( s ) > 0 .
[Equation]
where the last step follows from Lemma
[Equation]
Then, for any D f ≥ D KL , we have that
[Equation]
A. P ROOF OF T HEOREM
Proof.
We begin with
[Equation]
We have that
[Equation]
[Equation]
where the last step follows from recognizing that the inner-maximization is precisely the Fenchel conjugate of D f ( d ( s, a )
∥
d O ( s, a )) at R ( s ) + γ T V ( s, a ) − V ( s ) .
To show the relationship among V ⋆ and ξ ⋆ , we recognize that (27) and (9) are a pair of Fenchel primal-dual problems.
Lemma 3.
[Equation]
is the Fenchel dual to
[Equation]
Proof.
We define the indicator function δ X ( x ) as
[Equation]
∞ otherwise
Then, we define g : R |S| → R as g ( · ) :
=
δ { (1 − γ ) µ 0 } ( · ) .
Then, it can be shown that the Fenchel conjugate of g is g ⋆ ( · ) =
(1 − γ ) E µ 0 [ · ] .
In addition, we denote h ( · ) :
= D + f ( ·∥ d O ) ; then, h ⋆ ( · ) =
E ( s,a ) ∼ d
O
[ f ⋆ ( · )] .
Finally, define matrix operator A :
= γ T ⋆ − I .
Using these notations, we can write (27) as
[Equation]
Then, we proceed to derive the Fenchel dual of (30):
[Equation]
where (32) follows applying Fenchel conjugacy to h ⋆ , (34) follows from strong duality, (35) follows from the property of an adjoint operator, and (37) follows from applying Fenchel conjugacy to g ⋆ .
Here, we recognize that (37) is precisely the optimization problem (28)-(29), where we have moved the constraint (29) to the objective as the indicator function g ( − Ad ) :
[Equation]
Giving Lemma 3, we use the fact that d ∗ and V ∗ admit the following relationship:
[Equation]
This follows from the characterization of the optimal solutions for a pair of Fenchel primal-dual problems with convex g, h and linear operator A Nachum & Dai (2020).
In this case, assuming that we can exchange the order of expectation and derivative (e.g, conditions of Dominated Convergence
Theorem hold), we have
[Equation]
or equivalently,
[Equation]
as desired.
B R ELATED W ORK
Offline imitation learning.
The closest work is concurrent work, DEMODICE Kim et al.
(2022), a state-action based offline IL method, also using the DICE paradigm to estimates the occupancy ratio between the expert and the imitator; we overview the DICE literature in Appendix B.
Due to its dependence on expert actions, DEMODICE cannot be applied to the three problem settings we study.
At a technical level, a key limitation of DEMODICE is that it does not exploit the form of general
Fenchel duality and only support the KL-divergence, forgoing other f -divergences that can lead to more stable optimization Ghasemipour et al.
Ke et al.
Zhu et al.
Another related work is ORIL Zolna et al.
(2020), which adapts GAIL Ho & Ermon (2016) to the offline setting.
Finally, there has been recent work learning a pessimistic dynamics model using the offline dataset and then performs imitation learning by minimizing the state-action occupancy divergence with respect to the expert inside this learned model Chang et al.
As with DEMODICE, this approach requires expert actions and cannot be applied to the settings we study.
Imitation from observations, imitation with mismatched experts, and example-based RL All
three of these problems have been studied in the online setting.
IfO is often achieved through training an additional inverse dynamics model to infer the expert actions Torabi et al.
(2018; 2019); Liu et al.
; Radosavovic et al.
; in contrast, SMODICE matches the expert observations by identifying the correct actions supported in the offline data.
To handle experts with dynamics mismatch, some work explicitly learns a correspondence between the expert and the imitator MDPs Kim et al.
(2020); Raychaudhuri et al.
(2021); however, these approaches make much stronger assumptions on access to the expert MDP that are difficult to satisfy in the offline setting, such as demonstrations from auxillary tasks.
In contrast, SMODICE falls under the category of state-only imitation learning Liu et al.
; Radosavovic et al.
(2020), which overcomes expert dynamics differences by only matching the shared task-relevant state space (e.g., xy coordinates for locomotion tasks).
Finally, example-based RL was first studied in Eysenbach et al.
(2021); they introduce a recursive-classifier based off-policy actor critic method to solve it.
By casting this problem as state-occupancy matching between an imitator and a “teleporting” expert agent, SMODICE can solve the offline variant of this problem without modification.
Stationary distribution correction estimation.
Estimating the optimal policy’s stationary distribution using off-policy data was introduced by Nachum et al.
(2019a)
as the DICE trick.
This technique has been shown to be effective for off-policy evaluation Nachum et al. (2019a); Zhang* et al.
Dai et al.
policy optimization Nachum et al.
(2019b)
; Lee et al.
(2021), online imitation learning Kostrikov et al.
Zhu et al.
(2020), and concurrently, offline imitation learning Kim et al.
Within the subset of DICE-based policy optimization methods, none has tackled state-occupancy matching or directly apply Fenchel Duality to its full generality to arrive at the form of value function objective we derive.
C SMODICE WITH COMMON f - DIVERGENCES
Example 1 (SMODICE with χ 2 -divergence) .
Suppose f ( x ) =
( x − 1) 2 , corresponding to χ 2divergence.
Then, we can show that f ⋆ ( x ) = 12 ( x + 1) 2 and f ′ ⋆ ( x ) = x
+ 1 .
Hence, the SMODICE objective amounts to
[Equation]
[Equation]
and
[Equation]
Example 2 (SMODICE with KL-divergence) .
We have f ( x ) = x log x .
Using the fact that the conjugate of the negative entropy function, restricted to the probability simplex, is the log-sum-exp function Boyd et al.
(2004), it follows that D
⋆,f ( y ) =
log E
x ∼ q [exp y ( x )] .
Hence, the KL-divergence
SMODICE objective is
[Equation] [Equation]
and
[Equation] [Equation]
D SMODICE FOR T ABULAR MDP S
In this section, we derive the closed-form expression of SMODICE for tabular MDPs.
For simplicity, we assume that the expert state occupancies are given, d E ( s ) ∈ ∆( |S| ) .
A behavior policy π b is used to collect the offline dataset D O .
Then, we can construct a surrogate MDP ˆ M using maximum likelihood estimation (i.e., ˆ T ( s, a, s ′ )
=
n ( s,a,s ′ )
n ( s,a ) ).
Using ˆ M , we can extract the empirical estimate of the behavior policy occupancies d O ∈ ∆( |S||A| ) using linear programming.
Then, we can define the reward vector R ∈ R |S| + as R ( s ) =
log d E ( s )
d O ( s ) .
Using the χ 2 -divergence version of SMODICE, we can write down the objective for V ( s ) ∈ R |S| + :
[Equation]
[Equation]
We rewrite this expression in vector-matrix form to derive the closed-form solution.
To this end, we define T ∈ R |S||A|×|S| and B ∈ R |S||A|×|S| such that ( T V )( s, a ) =
(cid:80) s ′
T ( s ′ | s, a ) V ( s ′ ) and ( B V )( s, a ) =
V ( s ) .
Additionally, we denote µ 0 ∈ ∆( |S| ) and D = diag( d O ) ∈ R |S||A|×|S||A| .
Then, we can rewrite (45):
[Equation]
where r V ∈ R |S||A| and I is the all-one vector in R |S||A| .
Denoting J ( V ) :
= (1 − γ ) µ ⊤ 0 V + 12 ( r V +
I )
⊤ D ( r V + I )
[Equation]
it is clear that J ( V ) is a convex program in V .
Therefore
[Equation]
we can find its optimal solution by solving the first-order stationary point.
We have:
[Equation]
Then, by setting this expression to zero and solving for V gives the optimal V ∗ :
[Equation]
and we can recover ξ ∗ ( s, a ) =
d ∗ ( s,a ) d
O ( s,a ) :
[Equation]
[Equation]
Pythonic pseudo-code using NumPy Harris et al.
(2020) is given in Algorithm 1.
D. G RIDWORLD E XPERIMENTS
In this subsection, we provide more experimental details and analysis of the tabular SMODICE experiments shown in Figure 1.
To generate the offline dataset, a random policy (i.e., a policy that chooses each action with equal probabilities) is executed in the MDP for 10000 epsiodes.
We use this dataset to compute the approximate MDP.
Then, this MDP is used as an input to SMODICE (see Algorithm 1).
The data collection procedure for the offline imitation learning from examples setting is identical.
Offline IL from mismatched experts.
In this task, we consider an expert agent that can move one grid cell diagonally in any direction, whereas the imitator is only able to move one grid cell horizontally or vertically.
The expert policy is shown in black in Figure 2(a).
Using purely an offline dataset collected by a random agent, we compute the closed-form tabular SMODICE solution ( ?? )
using Algorithm 1 and obtain the zig-zagging policy shown in blue.
Indeed, this solution is one of the two correct solutions that minimize the state-occupancy divergence (the other one mirrors this path along the expert demo), while being feasible under the imitator dynamics.
Offline IL from examples.
We arbitrarily select a state to be the success state denoted by the green star in Figure 2(b).
In this case, the expert’s state occupancies is simply a one-hot vector with weight 1 at the success state.
Then, we again use the tabular version of SMODICE to compute the policy whose state occupancies is as close to this one-hot vector as possible; the solution is illustrated in blue.
As can be seen, this policy successfully reaches the goal.
Furthermore, it is easy to see that in this task, a policy that minimizes state-occupancy divergence to the expert (i.e., the one-hot vector) is one that reaches the goal with the fewest steps.
The policy learned by SMODICE is indeed among the set of optimal policies.
Furthermore, we compute the state occupancies of all states in the gridworld.
For the success state, d ( s ) ≈ 0 .
915 , whereas the second largest state occupancy is 0 .
This validates the intuition that (cid:80) [Equation]
E SMODICE WITH D EEP N EURAL N ETWORKS
For high-dimensional MDP with continuous state and action spaces, we instantiate SMODICE using deep neural networks.
In particular, we parameterize V θ and π ϕ using DNNs with weights θ and ϕ , respectively.
Remark.
We note that the sample-based estimation of Equation (16) (Line 9) is biased because T V is itself an expectation that is inside a (non-linear) convex function f Baird (1995); however, as in several prior works Nachum et al. (2019b); Nachum & Dai (2020); Lee et al.
(2021), we do not find this biased estimate to impact empirical performance and keep it for simplicity.
E. H YPERPARAMETERS AND A RCHITECTURE
We use the same hyperparameters for all SMODICE experiments in this paper modulo the choice of f -divergences (explained in the next section).
In terms of architecture, we use a simple 2-layer ReLU network with hidden size 256 to parameterize the value network.
For the policy network, we use the same architecture to parameterize a Gaussian output distribution; the mean and the log standard deviation are ouputs of two separate heads.
In addition, we use an tanh function on the Gaussian samples to enforce bounded actions, as in Haarnoja et al.
The discriminator uses the same architecture.
Table 1 summarizes the hyperparameters as well as the architecture.
E. C HOOSING f -D IVERGENCE IN P RACTICE
In our experiments, SMODICE is implemented using χ 2 -divergence for all tasks except Hopper,
Walker2d, and HalfCheetah.
Here, we show that a suitable choice of f -divergence can be chosen offline by observing the initial direction of the SMODICE policy loss on the offline dataset.
More specifically, on the environments in which SMODICE exhibited largest performance discrepancies between using KL-divergence or χ 2 -divergence, we have found that SMODICE returns are negatively correlated with the policy loss.
As shown in Figure 7, the poor performing variant of SMODICE
[Table]
always has a policy loss that initially jumps and vice-versa.
This makes intuitive sense given the composition of the offline datasets, which is a mix of small amount of expert data with a large amount of poor quality data (see Appendix G for more details).
When SMODICE fails to pick out the expert data, which is often narrowly distributed, then it must have assigned relatively higher importance weights to the lower quality data, which is more diverse.
This creates a more difficult supervised learning task, leading to higher training loss for the policy.
Therefore, in practice, we recommend monitoring SMODICE’s initial policy loss direction to determine whether the current f -divergence will lead to good performance and make changes accordingly.
F B ASELINES
TD3-BC.
Many of our baselines are implemented using TD3-BC as their offline policy optimizer.
We use the default hyperparameters for TD3-BC provided by Fujimoto & Gu (2021), shown in Table 2.
Implementation Details.
We use the official PyTorch implementation of TD3-BC, publicly available at https://github.com/sfujim/TD3_BC .
For DEMODICE, because the code is not public available, we implement it using PyTorch, adapting from https://github.com/secury/ optidice ; we use the hyperparameters reported in the paper.
Note that DEMODICE shares many architectures with SMODICE.
For example, DEMODICE uses a state-action discriminator, and we implement it by simply changing the input space of the state-based discriminator used in our SMODICE implementation.
For SAIL, we use the official PyTorch implementation ( https: //github.com/
FangchenLiu/SAIL ) and combine it with TD3-BC.
We implement RCE using PyTorch, adapting from the official TensorFlow implementation https://github.com/ google-research/google-research/tree/master/rce .
[Table]
G O FFLINE IL FROM O BSERVATIONS E XPERIMENTAL D ETAILS G. D ATASETS
For Hopper, Walker2d, HalfCheetah, Ant
[Table]
and AntMaze, we construct the offline datasets by combining a small amount of expert data and a large amount of low quality random data.
For the first four tasks, we leverage the respective “expert-v2” and “random-v2” datasets in the D4RL benchmark.
For AntMaze, we use trajectories from “antmaze-umaze-v2” as the expert data; for the random data, we simulate the antmaze environment for 1M steps using random actions and take the resulting transitions.
For the kitchen environment, we use the full “kitchen-mixed-v0” dataset as the offline dataset without further augmentation.
See Table 3 for dataset breakdown.
In AntMaze (Figure 4(b)), an Ant agent is tasked with navigating an U-shaped maze from one end to the other end (i.e., the goal region).
The offline dataset (i.e., “antmaze-umaze-v2”) consists of trajectories ( ≈ 300k transitions) of an Ant agent navigating to the goal region from initial states;
The trajectories are not always successful; often, the Ant flips over to its legs before it reaches the goal.
As above, we additionally include 1 million random-action transitions to increase the task difficulty.
We take one trajectory from the offline dataset that successfully reaches the goal to be the expert trajectory.
Franka Kitchen (Figure 4(c)), introduced by Gupta et al.
(2019), involves controlling a 9-DoF Franka robot to manipulate common household kitchen objects (e.g., microwave, kettle, cabinet) sequentially to achieve a pre-specified configuration of objects.
The dataset (i.e., “kitchen-mixed-v0”) consists of undirected human teleoperated demonstrations, meaning that each trajectory only solves a subset of the tasks.
Together, these six tasks (illustrated in Figure 5.1) require scalability to high-dimensional state-action spaces and robustness to different dataset compositions.
[Table]
G. A DDITIONAL R ESULTS
In this section, we present some additional results as well as ablation experiments.
Diverse AntMaze.
In Section 5.1, we have found that two of the baselines (SAIL-TD3-BC and
ORIL) outperform SMODICE on the AntMaze benchmark.
To investigate their sources of empirical gain, we have designed a diverse version of the AntMaze dataset to test how different approaches are robust to the dataset composition on the same task.
To this end, we take the AntMaze offline dataset (explained above) and reverse half of the trajectories in their directions.
In other words, these reversed trajectories would navigate from the original goal to the initial state.
This procedure is easy to do because the U-shaped maze is symmetric.
Then, using this dataset, we have trained all approaches in Section 5.1 again.
As shown in Figure 9(a), on this dataset, both SAIL-TD3-BC and
ORIL quickly collapse, indicating that these methods are very brittle to the dataset composition.
In contrast, SMODICE remains the best performing algorithm, despite overall drop in all methods’ performances.
SMODICE with Zero Reward.
We compare SMODICE with SMODICE-Zero, which simply assigns every transition zero reward (i.e., R ( s ) = 0 ) regardless of its similarity to an expert state.
Then, we compare the ratio of the importance weights (i.e., ξ ( s, a ) ) assigned to the offline expert data and the offline random data by the two SMODICE methods, respectively.
As shown in Figure 8,
SMODICE assigns much higher relative weights to the expert data and consequently significantly outperforms SMODICE-Zero.
These results demonstrate that SMODICE’s empirical performance comes from its superior ability to discriminate the offline expert data, which is a by-product of its optimization procedure.
[Figure]
H O FFLINE IL FROM MISMATCHED E XPERT E XPERIMENTAL D ETAILS H. C ONTINUOUS C ONTROL E XPERIMENTS
Datasets and Baselines.
We compare SMODICE to SAIL-TD3-BC and ORIL, which are both state-based offline IL methods; in particular, we note that SAIL is originally designed to be robust to mismatched experts.
We consider only tasks in which both SAIL-TD3-BC and ORIL obtained non-trivial performance, including HalfCheetah, Ant, and AntMaze.
Then, for each environment, we train a mismatched expert and collect one expert trajectory, replacing the original expert trajectory used in Section 5.1.
The mismatched experts for the respective tasks are (i) “HalfCheetah-Short”, where the torso of the cheetah agent is halved in length, (ii) “Ant-Disabled”, where the front legs are shrank by a quarter in length, and (iii) a 2D PointMass agent operating in the same maze configuration.
The mismatched experts are illustrated in Figure 10 in Appendix H. For the first two, we train an expert policy using SAC Haarnoja et al.
and collect one expert trajectory.
The latter task is already in D4RL; thus, we take one trajectory from “maze2d-umaze-v0” as the expert trajectory.
Because Ant and PointMass have different state spaces, following Liu et al.
(2019), we train the discriminator on the shared xy -coordinates of the two state spaces.
The offline datasets are identical to the ones in Section 5.1.
Mismatched Experts.
The mismatched experts are illustrated in Figure 10.
Comparison between PointMass and Ant experts for AntMaze.
The trajectories of PointMass and Ant experts are illustrated in Figure 11.
As can be seen, the PointMass trajectory is more regular and smooth due to its simpler dynamics and the use of a waypoint controller.
In contrast, the ant trajectory is much less well-behaved because solving the maze task using the Ant agent is intrinsically a difficult task; consequently, it is difficult to provide an Ant demonstration.
This example serves as a strong motivating problem for offline imitation learning with mismatched experts.
[Figure]
[Table]
H. Q UANTITATIVE A NALYSIS OF F IGURE
We quantitatively measure the percentage drop-in-performance for each method in Figure 5, computed as | max original − max mismatched | max original .
Note that this metric favors the baselines as taking the maximum value advantages methods that are more unstable.
Nevertheless, as shown in Table 4, SMODICE is still by far the most robust method overall and in each individual task.
As expected, ORIL does the worst as it is not designed to handle mismatched dynamics; this shows that using a state-based discriminator in itself is not sufficient.
I O FFLINE IL FROM E XAMPLES E XPERIMENTAL D ETAILS I. D ATASETS
We collect 300 success-state examples for each of the tasks.
The examples are randomly sampled from the subset of the offline dataset that achieves the task.
Task success is verfied through a pre-defined sparse reward function (e.g., distance threshold function).
I. E NVIRONMENTS .
PointMass-4Direction.
This environment is adapted from the “maze2d-umaze-v0” environment in
D4RL by changing the map configuration.
The environment termination condition is triggered when the agent successfully comes within a small radius of the specified goal.
AntMaze-Example.
This environment is identical to the environments used in previous two settings.
Kettle and Microwave.
These environments are adapted from the ”kitchen-mixed-v0” environment in D4RL.
The environments are identical as the original except the termination conditions.
Both of these tasks terminate when the Franka robot places the specified object within a small radius of the desired configuration.
I. E XAMPLES OF S UCCESS S TATES
All success states are extracted from the offline dataset used for policy training.
We illustrate one representative example from each task in Figure 13.
[Figure]
I. A DDITIONAL A NALYSIS
To better understand SMODICE, on PointMass-4Direction, we visualize the importance weights ξ ( s, a ) it assigns to the offline dataset.
As shown in Figure 12, SMODICE assigns much higher weights to transitions along the correct path from the initial state region to the success examples.
Interestingly, the weights progressively decrease along this path, indicating that SMODICE has learned that it must pay more attention transitions at the beginning of the path, since making a mistake there is more likely to derail progress towards the goal.
This behavior occurs automatically via
SMODICE’s state-matching objective without any additional bias.
# d_E: the expert state occupancies, |S| # mdp: the empirical MDP learned using offline data # pi_b: the behavior policy, |S||A| def SMODICE(mdp, d_E, pi_b): d_O_sa = compute_policy_occupancies(mdp, pi_b) # |S||A| d_O =
d_O_sa.reshape(mdp.
S, mdp.
A).sum(axis=1)
# |S| # compute reward function
R = np.log(d_E/d_O)
# |S| # define and reshape matrices
T = mdp.
T.reshape(mdp.
S * mdp.
A, mdp.S)
# |S||A| x |S|
B = np.repeat(np.eye(mdp.
S), mdp.
A, axis=0) # |S||A| x |S|
I = np.ones(mdp.
S * mdp.
A)
# |S||A|
D = np.diag(d_O_sa) # |S||A| x |S||A| # compute optimal V
H =
(mdp.gamma * P - B).T @
D @
(mdp.gamma * T - B) # |S| x |S| y = -((1 - mdp.gamma) * p0 + (mdp.gamma * P - B).T @
D @
(I + B @ R))
# |S|
V_star = np.linalg.pinv(H) @
y # |S| # compute optimal occupancy ratios xi_star =
B @ R + (mdp.gamma * P - B) @
V_star + 1 # |S||A| m =
np.array(xi_star >= 0, dtype=np.float)
xi_star
=
xi_star
* m
1: Require : Expert demonstration(s) D E , offline dataset D O , choice of f -divergence f 2: Randomly initialize discriminator c ψ , value function V θ , and policy π ϕ .
Train Expert (resp.
Example) Discriminator 4: Train c ψ using D E and D O using Equation (15) 5: //
Train Lagrangian Value Function 6: for number of iterations do 7:
Sample minibatch of offline data { s it , a it , s it +1 } Ni =1 ∼ D O , { s i 0 } Mi =1 ∼ D O ( µ 0 ) 8:
Obtain reward: R i =
c θ ( s it ) , i = 1 , ..., N 9:
Compute value objective L ( θ ) :
= (1 − γ ) 1 M (cid:80) Mi
=1 V θ
( s i 0 )
+ 1 N f ⋆ (cid:0) R i + γV
( s it +1 ) − V
( s it ) (cid:1) 10:
Update V θ using SGD: V θ ← V θ − η V ∇L ( θ ) 11: end for 12: //
Policy Learning 13: for number of iterations do 14:
Sample minibatch of offline data { s it , a it , s it +1 } Ni =1 ∼ D O 15: //
Compute Optimal Importance Weights 16:
Compute ξ ∗ ( s i , a i ) = f ′
⋆ (cid:0) R ( s i ) + γV ( s it +1 )
− V ( s it ) (cid:1) , i = 1 , ..., N 17: //
Weighted Behavior Cloning 18:
Update π ψ using Equation (18) 19: end for