SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jinxin Liu $^{1,2*}$  Li He $^{1*}$  Yachen Kang $^{1,2}$  Zifeng Zhuang $^{1,2}$
Donglin Wang $^{1,4}$  Huazhe Xu $^{3,5,6}$
$^{1}$ Westlake University  $^{2}$ Zhejiang University  $^{3}$ Tsinghua University  $^{4}$ Westlake Institute for Advanced Study  $^{5}$ Shanghai Qi Zhi Institute  $^{6}$ Shanghai AI Lab
In this paper, we present ContExtual Imitation Learning (CEIL), a general and broadly applicable algorithm for imitation learning (IL). Inspired by the formulation of hindsight information matching, we derive CEIL by explicitly learning a hindsight embedding function together with a contextual policy using the hindsight embeddings. To achieve the expert matching objective for IL, we advocate for optimizing a contextual variable such that it biases the contextual policy towards mimicking expert behaviors. Beyond the typical learning from demonstrations (LfD) setting, CEIL is a generalist that can be effectively applied to multiple settings including: 1) learning from observations (LfO), 2) offline IL, 3) cross-domain IL (mismatched experts), and 4) one-shot IL settings. Empirically, we evaluate CEIL on the popular MuJoCo tasks (online) and the D4RL dataset (offline). Compared to prior state-of-the-art baselines, we show that CEIL is more sample-efficient in most online IL tasks and achieves better or competitive performances in offline tasks.
Imitation learning (IL) allows agents to learn from expert demonstrations. Initially developed with a supervised learning paradigm [58, 63], IL can be extended and reformulated with a general expert matching objective, which aims to generate policies that produce trajectories with low distributional distances to expert demonstrations [30]. This formulation allows IL to be extended to various new settings: 1) online IL where interactions with the environment are allowed, 2) learning from observations (LfO) where expert actions are absent, 3) offline IL where agents learn from limited expert data and a fixed dataset of sub-optimal and reward-free experience, 4) cross-domain IL where the expert demonstrations come from another domain (i.e., environment) that has different transition dynamics, and 5) one-shot IL which expects to recover the expert behaviors when only one expert trajectory is observed for a new IL task.
Modern IL algorithms introduce various designs or mathematical principles to cater to the expert matching objective in a specific scenario. For example, the LfO setting requires particular considerations regarding the absent expert actions, e.g., learning an inverse dynamics function [5, 65]. Besides, out-of-distribution issues in offline IL require specialized modifications to the learning objective, such as introducing additional policy/value regularization [32, 72]. However, such a methodology, designing an individual formulation for each IL setting, makes it difficult to scale up a specific IL algorithm to more complex tasks beyond its original IL setting, e.g., online IL methods often suffer severe performance degradation in offline IL settings. Furthermore, realistic IL tasks are often not subject to a particular IL setting but consist of a mixture of them. For example, we may have access
*Equal contributions. Corresponding author: Donglin Wang <wangdonglin@westlake.edu.cn>
Table 1: A coarse summary of IL methods demonstrating 1) different expert data modalities they can handle (learning from demonstrations or observations), 2) disparate task settings they consider (learning from online environment interactions or pre-collected offline static dataset), 3) the specific cross-domain setting they assume (the transition dynamics between the learning environment and that of the expert behaviors are different), and 4) the unique one-shot merit they desire (the learned policy is capable of one-shot transfer to new imitation tasks). We highlight that our contextual imitation learning (CEIL) method can naturally be applied to all the above IL settings.
Expert data	Task setting	Cross-domain	One-shot
LfD	LfO	Online	Offline
S-on-LfD [9, 13, 21, 30, 38, 52, 57, 61, 77]	✓	X	✓	X	X	X
S-on-LfO [7, 54, 65, 66, 75]	✓	✓	✓	X	X	X
S-off-LfD [19, 32, 33, 39, 55, 70, 72, 73]	✓	X	X	✓	X	X
S-off-LfO [78]	✓	✓	X	✓	X	X
C-on-LfD [18, 69, 79]	✓	X	✓	X	✓	X
C-on-LfO [20, 25, 26, 48, 59, 60]	✓	✓	✓	X	✓	X
C-off-LfD [34]	✓	X	X	✓	✓	X
C-off-LfO [56, 68]	✓	✓	X	✓	✓	X
S-on/off-LfO [28]	✓	✓	✓	✓	X	X
Online one-shot [14, 16, 40]	✓	X	✓	X	X	✓
Offline one-shot [24, 71]	✓	X	X	✓	X	✓
CEIL (ours)	✓	✓	✓	✓	✓	✓
to both demonstrations and observation-only data in offline robot tasks; however, it could require significant effort to adapt several specialized methods to leverage such mixed/hybrid data. Hence, a problem naturally arises: How can we accommodate various design requirements of different IL settings with a general and practically ready-to-deploy IL formulation?
Hindsight information matching, a task-relabeling paradigm in reinforcement learning (RL), views control tasks as analogous to a general sequence modeling problem, with the goal to produce a sequence of actions that induces high returns [12]. Its generality and simplicity enable it to be extended to both online and offline settings [17, 42]. In its original RL context, an agent directly uses known extrinsic rewards to bias the hindsight information towards task-related behaviors. However, when we attempt to retain its generality in IL tasks, how to bias the hindsight towards expert behaviors remains a significant barrier as the extrinsic rewards are missing.
To design a general IL formulation and tackle the above problems, we propose ContExtual Imitation Learning (CEIL), which readily incorporates the hindsight information matching principle within a bi-level expert matching objective. In the inner-level optimization, we explicitly learn a hindsight embedding function to deal with the challenges of unknown rewards. In the outer-level optimization, we perform IL expert matching via inferring an optimal embedding (i.e., hindsight embedding biasing), replacing the naive reward biasing in hindsight. Intuitively, we find that such a bi-level objective results in a spectrum of expert matching objectives from the embedding space to the trajectory space. To shed light on the applicability and generality of CEIL, we instantiate CEIL to various IL settings, including online/offline IL, LfD/LfO, cross-domain IL, and one-shot IL settings.
In summary, this paper makes the following contributions: 1) We propose a bi-level expert matching objective ContExtual Imitation Learning (CEIL), inheriting the spirit of hindsight information matching, which decouples the learning policy into a contextual policy and an optimal embedding. 2) CEIL exhibits high generality and adaptability and can be instantiated over a range of IL tasks. 3) Empirically, we conduct extensive empirical analyses showing that CEIL is more sample-efficient in online IL and achieves better or competitive results in offline IL tasks.
Recent advances in decision-making have led to rapid progress in IL settings (Table 1), from typical learning from demonstrations (LfD) to learning from observations (LfO) [7, 9, 35, 54, 62, 66], from online IL to offline IL [11, 15, 33, 53, 73], and from single-domain IL to cross-domain IL [34, 48, 56, 68]. Targeting a specific IL setting, individual works have shown their impressive
ability to solve the exact IL setting. However, it is hard to retrain their performance in new unprepared IL settings. In light of this, it is tempting to consider how we can design a general and broadly applicable IL method. Indeed, a number of prior works have studied part of the above IL settings, such as offline LfO [78], cross-domain LfO [48, 60], and cross-domain offline IL [56]. While such works demonstrate the feasibility of tackling multiple IL settings, they still rely on standard online/offline RL algorithmic advances to improve performance [25, 32, 44, 47, 50, 51, 55, 72, 76]. Our objective diverges from these works, as we strive to minimize the reliance on the RL pipeline by replacing it with a simple supervision objective, thus avoiding the dependence on the choice of RL algorithms.
Our approach to IL is most closely related to prior hindsight information-matching methods [2, 8, 24, 49], both learning a contextual policy and using a contextual variable to guide policy improvement. However, these prior methods typically require additional mechanisms to work well, such as extrinsic rewards in online RL [4, 42, 64] or a handcrafted target return in offline RL [12, 17]. Our method does not require explicit handling of these components. By explicitly learning an embedding space for both expert and suboptimal behaviors, we can bias the contextual policy with an inferred optimal embedding (contextual variable), thus avoiding the need for explicit reward biasing in prior works. Our method also differs from most prior offline transformer-based RL/IL algorithms that explicitly model a long sequence of transitions [10, 12, 31, 36, 43, 71]. We find that simple fully-connected networks can also elicit useful embeddings and guide expert behaviors when conditioned on a well-calibrated embedding. In the context of the recently proposed prompt-tuning paradigm in large language tasks or multi-modal tasks [27, 45, 74], our method can be interpreted as a combination of IL and prompting-tuning, with the main motivation that we tune the prompt (the optimal contextual variable) with an expert matching objective in IL settings.
Before discussing our method, we briefly introduce the background for IL, including learning from demonstrations (LfD), learning from observations (LfO), online IL, offline IL, and cross-domain settings in Section 3.1, and introduce the hindsight information matching in Section 3.2.
3.1 Imitation Learning
Consider a control task formulated as a discrete-time Markov decision process  $(\mathrm{MDP})^2\mathcal{M} = \{\mathcal{S},\mathcal{A},\mathcal{T},r,\gamma ,p_0\}$ , where  $\mathcal{S}$  is the state (observation) space,  $\mathcal{A}$  is the action space,  $\mathcal{T}:S\times \mathcal{A}\times \mathcal{S}\to \mathbb{R}$  is the transition dynamics function,  $r:S\times \mathcal{A}\rightarrow \mathbb{R}$  is the reward function,  $\gamma$  is the discount factor, and  $p_0$  is the distribution of initial states. The goal in a reinforcement learning (RL) control task is to learn a policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s})$  maximizing the expected sum of discounted rewards  $\mathbb{E}_{\pi_{\theta}(\boldsymbol {\tau})}\left[\sum_{t = 0}^{T - 1}\gamma^t r(\mathbf{s}_t,\mathbf{a}_t)\right]$ , where  $\boldsymbol {\tau}:= \{\mathbf{s}_0,\mathbf{a}_0,\dots ,\mathbf{s}_{T - 1},\mathbf{a}_{T - 1}\}$  denotes the trajectory and the generated trajectory distribution  $\pi_{\theta}(\boldsymbol {\tau}) = p_0(\mathbf{s}_0)\pi_{\theta}(\mathbf{a}_0|\mathbf{s}_0)\prod_{t = 1}^{T - 1}\pi_{\theta}(\mathbf{a}_t|\mathbf{s}_t)\mathcal{T}(\mathbf{s}_t|\mathbf{s}_{t - 1},\mathbf{a}_{t - 1})$ .
In IL, the ground truth reward function (i.e.,  $r$  in  $\mathcal{M}$ ) is not observed. Instead, we have access to a set of demonstrations (or observations)  $\{\pmb{\tau} | \pmb{\tau} \sim \pi_E(\pmb{\tau})\}$  that are collected by an unknown expert policy  $\pi_E(\mathbf{a}|\mathbf{s})$ . The goal of IL tasks is to recover a policy that matches the corresponding expert policy. From the mathematical perspective, IL achieves the plain expert matching objective by minimizing the divergence of trajectory distributions between the learner and the expert:
$$
\min  _ {\pi_ {\theta}} D \left(\pi_ {\theta} (\boldsymbol {\tau}), \pi_ {E} (\boldsymbol {\tau})\right), \tag {1}
$$
where  $D$  is a distance measure. Meanwhile, we emphasize that the given expert data  $\{\tau|\tau \sim \pi_E(\tau)\}$  may not contain the corresponding expert actions. Thus, in this work, we consider two IL cases where the given expert data  $\tau$  consists of a set of state-action demonstrations  $\{(\mathbf{s}_t,\mathbf{a}_t,\mathbf{s}_{t + 1})\}$  (learning from demonstrations, LfD), as well as a set of state-only transitions  $\{(\mathbf{s}_t,\mathbf{s}_{t + 1})\}$  (learning from observations, LfO). When it is clear from context, we abuse notation  $\pi_E(\tau)$  to denote both demonstrations in LfD and observations in LfO for simplicity.
Besides, we can also divide IL settings into two orthogonal categories: online IL and offline IL. In online IL, the learning policy  $\pi_{\theta}$  can interact with the environment and generate online trajectories  $\pmb{\tau} \sim \pi_{\theta}(\pmb{\tau})$ . In offline IL, the agent cannot interact with the environment but has access to an offline
2In this paper, we use environment and MDP interchangeably, and use state and observation interchangeably.
static dataset  $\{\pmb {\tau}|\pmb {\tau}\sim \pi_{\beta}(\pmb {\tau})\}$ , collected by some unknown (sub-optimal) behavior policies  $\pi_{\beta}$ . By leveraging the offline data  $\{\pi_{\beta}(\pmb {\tau})\} \cup \{\pi_{E}(\pmb {\tau})\}$  without any interactions with the environment, the goal of offline IL is to learn a policy recovering the expert behaviors (demonstrations or observations) generated by  $\pi_{E}$ . Note that, in contrast to the typical offline RL problem [46], the offline data  $\{\pi_{\beta}(\pmb {\tau})\}$  in offline IL does not contain any reward signal.
Cross-domain IL. Beyond the above two IL branches (online/offline and LfD/LfO), we can also divide IL into: 1) single-domain IL and 2) cross-domain IL, where 1) the single-domain IL assumes that the expert behaviors are collected in the same MDP in which the learning policy is to be learned, and 2) the cross-domain IL studies how to imitate expert behaviors when discrepancies exist between the expert and the learning MDPs (e.g., differing in their transition dynamics or morphologies).
3.2 Hindsight Information Matching
In typical goal-conditioned RL problems, hindsight experience replay (HER) [3] proposes to leverage the rich repository of the failed experiences by replacing the desired (true) goals of training trajectories with the achieved goals of the failed experiences:
$$
\operatorname {A l g} \left(\pi_ {\theta}; \boldsymbol {g}, \tau_ {\boldsymbol {g}}\right)\rightarrow \operatorname {A l g} \left(\pi_ {\theta}; f _ {\text {H E R}} \left(\tau_ {\boldsymbol {g}}\right), \tau_ {\boldsymbol {g}}\right),
$$
where the learner  $\mathrm{Alg}(\pi_{\theta};\cdot ,\cdot)$  could be any RL methods,  $\tau_g\sim \pi_\theta (\tau_g|g)$  denotes the trajectory generated by a goal-conditioned policy  $\pi_{\theta}(\mathbf{a}_{t}|s_{t},g)$ , and  $f_{\mathrm{HER}}$  denotes a pre-defined (hindindsight information extraction) function, e.g., returning the last state in trajectory  $\tau_{g}$ .
HER can also be applied to the (single-goal) reward-driven online/offline RL tasks, setting the return (sum of the discounted rewards) of a trajectory as an implicit goal for the corresponding trajectory. Thus, we can reformulate the (single-goal) reward-driven RL task, learning policy  $\pi_{\theta}(\mathbf{a}_t|\mathbf{s}_t)$  that maximize the return, as a multi-goal RL task, learning a return-conditioned policy  $\pi_{\theta}(\mathbf{a}_t|\mathbf{s}_t,\cdot)$  that maximize the following log-likelihood:
$$
\max  _ {\pi_ {\theta}} \mathbb {E} _ {\mathcal {D} (\boldsymbol {\tau})} \left[ \log \pi_ {\theta} \left(\mathbf {a} | \mathbf {s}, f _ {\mathrm {R}} (\boldsymbol {\tau})\right) \right], \tag {2}
$$
where  $f_{\mathrm{R}}(\tau)$  denotes the return of trajectory  $\pmb{\tau}$ . At test, we can then condition the contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\cdot)$  on a desired target return. In offline RL, the empirical distribution  $\mathcal{D}(\tau)$  in Equation 2 can be naturally set as the offline data distribution; in online RL,  $\mathcal{D}(\tau)$  can be set as the replay/experience buffer, and will be updated and biased towards trajectories that have high expected returns.
Intuitively, biasing the sampling distribution  $(\mathcal{D}(\tau)$  towards higher returns) leads to an implicit policy improvement operation. However, such an operator is non-trivial to obtain in the IL problem, where we do not have access to a pre-defined function  $f_{\mathrm{R}}(\tau)$  to bias the learning policy towards recovering the given expert data  $\{\pi_E(\tau)\}$  (demonstrations or observations).
In this section, we will formulate IL as a bi-level optimization problem, which will allow us to derive our method, contextual imitation learning (CEIL). Instead of attempting to train the learning policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s})$  with the plain expert matching objective (Equation 1), our approach introduces an additional contextual variable  $\mathbf{z}$  for a contextual IL policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\cdot)$ . The main idea of CEIL is to learn a contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z})$  and an optimal contextual variable  $\mathbf{z}^*$  such that the given expert data (demonstrations in LfD or observations in LfO) can be recovered by the learned  $\mathbf{z}^*$ -conditioned policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z}^*)$ . We begin by describing the overall framework of CEIL in Section 4.1, and make a connection between CEIL and the plain expert matching objective in Section 4.2, which leads to a practical implementation under various IL settings in Section 4.3.
4.1 Contextual Imitation Learning (CEIL)
Motivated by the hindsight information matching in online/offline RL (Section 3.2), we propose to learn a general hindsight embedding function  $f_{\phi}$ , which encodes trajectory  $\pmb{\tau}$  (with window size  $T$ ) into a latent variable  $\mathbf{z} \in \mathcal{Z}$ ,  $|\mathcal{Z}| \ll T * |\mathcal{S}|$ . Formally, we learn the embedding function  $f_{\phi}$  and a corresponding contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z})$  by minimizing the trajectory self-consistency loss:
$$
\pi_ {\theta}, f _ {\phi} = \min  _ {\pi_ {\theta}, f _ {\phi}} - \mathbb {E} _ {\mathcal {D} (\boldsymbol {\tau})} [ \log \pi_ {\theta} (\boldsymbol {\tau} | f _ {\phi} (\boldsymbol {\tau})) ] = \min  _ {\pi_ {\theta}, f _ {\phi}} - \mathbb {E} _ {\boldsymbol {\tau} \sim \mathcal {D} (\boldsymbol {\tau})} \mathbb {E} _ {(\mathbf {s}, \mathbf {a}) \sim \boldsymbol {\tau}} [ \log \pi_ {\theta} (\mathbf {a} | \mathbf {s}, f _ {\phi} (\boldsymbol {\tau})) ], \tag {3}
$$
where in the online setting, we sample trajectory  $\tau$  from buffer  $\mathcal{D}(\tau)$ , known as the experience replay buffer in online RL; in the offline setting, we sample trajectory  $\tau$  directly from the given offline data.
If we can ensure that the learned contextual policy  $\pi_{\theta}$  and the embedding function  $f_{\phi}$  are accurate on the empirical data  $\mathcal{D}(\tau)$ , then we can convert the IL policy optimization objective (in Equation 1) into a bi-level expert matching objective:
$$
\min  _ {\mathbf {z} ^ {*}} D \left(\pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*}), \pi_ {E} (\boldsymbol {\tau})\right), \tag {4}
$$
$$
\text {s . t .} \pi_ {\theta}, f _ {\phi} = \min  _ {\pi_ {\theta}, f _ {\phi}} - \mathbb {E} _ {\mathcal {D} (\boldsymbol {\tau})} \left[ \log \pi_ {\theta} (\boldsymbol {\tau} | f _ {\phi} (\boldsymbol {\tau})) \right] - \mathcal {R} (f _ {\phi}), \text {a n d} \mathbf {z} ^ {*} \in f _ {\phi} \circ \operatorname {s u p p} (\mathcal {D}), \tag {5}
$$
where  $\mathcal{R}(f_{\phi})$  is an added regularization over the embedding function (we will elaborate on it later), and  $\mathrm{supp}(\mathcal{D})$  denotes the support of the trajectory distribution  $\{\tau | \mathcal{D}(\tau) > 0\}$ . Here  $f_{\phi}$  is employed to map the trajectory space to the latent variable space  $(\mathcal{Z})$ . Intuitively, by optimizing Equation 4, we expect the induced trajectory distribution of the learned  $\pi_{\theta}(\mathbf{a} | \mathbf{s}, \mathbf{z}^{*})$  will match that of the expert. However, in the offline IL setting, the contextual policy can not interact with the environment. If we directly optimize the expert matching objective (Equation 4), such an objective can easily exploit generalization errors in the contextual policy model to infer a mistakenly overestimated  $\mathbf{z}^{*}$  that achieves low expert-matching loss but does not preserve the trajectory self-consistency (Equation 3). Therefore, we formalize CEIL into a bi-level optimization problem, where, in Equation 5, we explicitly constrain the inferred  $\mathbf{z}^{*}$  lies in the  $(f_{\phi}$ -mapped) support of the training trajectory distribution.
At a high level, CEIL decouples the learning policy into two parts: an expressive contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\cdot)$  and an optimal contextual variable  $\mathbf{z}^*$ . By comparing CEIL with the plain expert matching objective,  $\min_{\pi_{\theta}}D(\pi_{\theta}(\pmb {\tau}),\pi_{E}(\pmb {\tau}))$ , in Equation 1, we highlight two merits: 1) CEIL's expert matching loss (Equation 4) does not account for updating  $\pi_{\theta}$  and is only incentivized to update the low-dimensional latent variable  $\mathbf{z}^*$ , which enjoys efficient parameter learning similar to the prompt tuning in large language models [74], and 2) we learn  $\pi_{\theta}$  by simply performing supervised regression (Equation 5), which is more stable compared to vanilla inverse-RL/adversarial-IL methods.
4.2 Connection to the Plain Expert Matching Objective
To gain more insight into Equation 4 that captures the quality of IL (the degree of similarity to the expert data), we define  $D(\cdot, \cdot)$  as the sum of reverse KL and forward KL divergence<sup>3</sup>, i.e.,  $D(q, p) = D_{\mathrm{KL}}(q \| p) + D_{\mathrm{KL}}(p \| q)$ , and derive an alternative form for Equation 4:
$$
\arg \min  _ {\mathbf {z} ^ {*}} D \left(\pi_ {\theta} \left(\boldsymbol {\tau} \mid \mathbf {z} ^ {*}\right), \pi_ {E} (\boldsymbol {\tau})\right) = \arg \max  _ {\mathbf {z} ^ {*}} \underbrace {\mathcal {I} \left(\mathbf {z} ^ {*} ; \boldsymbol {\tau} _ {E}\right) - \mathcal {I} \left(\mathbf {z} ^ {*} ; \boldsymbol {\tau} _ {\theta}\right)} _ {\mathcal {I} _ {\mathrm {M I}}} - \underbrace {D \left(\pi_ {\theta} (\boldsymbol {\tau}) , \pi_ {E} (\boldsymbol {\tau})\right)} _ {\mathcal {I} _ {D}}, \tag {6}
$$
where  $\mathcal{I}(\mathbf{x};\mathbf{y})$  denotes the mutual information (MI) between  $\mathbf{x}$  and  $\mathbf{y}$ , which measures the predictive power of  $\mathbf{y}$  on  $\mathbf{x}$  (or vice-versa), the latent variables are defined as  $\pmb{\tau}_{E} := \pmb{\tau} \sim \pi_{E}(\pmb{\tau})$ ,  $\pmb{\tau}_{\theta} := \pmb{\tau} \sim p(\mathbf{z}^{*})\pi_{\theta}(\pmb{\tau}|\mathbf{z}^{*})$ , and  $\pi_{\theta}(\pmb{\tau}) = \mathbb{E}_{\mathbf{z}^{*}}[\pi_{\theta}(\pmb{\tau}|\mathbf{z}^{*})]$ .
Intuitively, the second term  $\mathcal{I}_D$  on RHS of Equation 6 is similar to the plain expert matching objective in Equation 1, except that here we optimize a latent variable  $\mathbf{z}^*$  over this objective. Regarding the MI terms  $\mathcal{I}_{\mathrm{MI}}$ , we can interpret the maximization over  $\mathcal{I}_{\mathrm{MI}}$  as an implicit policy improvement, which incentivizes the optimal latent variable  $\mathbf{z}^*$  for having high predictive power of the expert data  $\tau_{E}$  and having low predictive power of the non-expert data  $\tau_{\theta}$ .
Further, we can rewrite the MI term  $(\mathcal{I}_{\mathrm{MI}}$  in Equation 6) in terms of the learned embedding function  $f_{\phi}$ , yielding an approximate embedding inference objective  $\mathcal{I}_{\mathrm{MI}(f_{\phi})}$ :