PWIL builds imitation rewards by transporting an agent’s state–action visits toward expert demonstrations. Each expert sample has limited capacity, so repeatedly visiting one convenient match cannot absorb the whole episode’s mass. The transport is greedy and the reward nonlinear; neither operation is exact global Wasserstein minimization.
Empirical measures
An episode supplies \(x_i=(s_i,a_i)\in\mathbb R^d\), \(i=1,\ldots,T\); demonstrations supply \(y_j\), \(j=1,\ldots,D\). Their empirical probability measures have total mass one:
\[\widehat\rho_\pi=\frac1T\sum_{i=1}^{T}\delta_{x_i},\qquad \widehat\rho_E=\frac1D\sum_{j=1}^{D}\delta_{y_j}.\]They retain visit frequencies but discard temporal order. Identical multisets of state–action pairs therefore produce identical measures; timing constraints must enter through the environment, policy, or representation.
Transport couplings
For metric space \((M,d_M)\), \(p\ge1\), and measures with finite \(p\)th moments,
\[W_p(\mu,\nu)^p =\inf_{\gamma\in\Pi(\mu,\nu)} \int_{M\times M}d_M(x,y)^p\,d\gamma(x,y).\]The couplings \(\Pi(\mu,\nu)\) have marginals \(\mu\) and \(\nu\). For uniform empirical measures, the nonnegative matrix \(\Gamma\in\mathbb R_+^{T\times D}\) satisfies
\[\sum_{j=1}^{D}\Gamma_{ij}=\frac1T, \qquad \sum_{i=1}^{T}\Gamma_{ij}=\frac1D.\]These probability marginals are not row and column sums of one, which would be incompatible when \(T\ne D\). The discrete problem is
\[W_p^p(\widehat\rho_\pi,\widehat\rho_E) =\min_{\Gamma\in\Pi}\sum_{i,j}\Gamma_{ij}d_M(x_i,y_j)^p.\]and the one-point identity \(W_p(\delta_x,\delta_y)=d_M(x,y)\) checks power and normalization.
PWIL’s episodic objective is
\[\min_\pi\;\mathbb E_{\tau\sim\pi} \left[W_1(\widehat\rho_\pi(\tau),\widehat\rho_E)\right].\]It need not equal transport between infinite-horizon occupancy measures: expectation and transport cannot freely be exchanged. Reward construction avoids jointly fitting an adversarial discriminator (Dadashi et al.). For an exact completed-episode coupling, row costs
\[c_i^*=\sum_j\Gamma^*_{ij}d_M(x_i,y_j).\]satisfy \(\sum_i c_i^*=W_1\), but depend on future states unavailable online.
Greedy allocation
An episode starts with expert capacities \(b_j=1/D\). At step \(i\), source mass \(1/T\) is assigned to available samples in increasing distance order:
\[\min_{g\in\mathbb R_+^D}\sum_jg_jd_M(x_i,y_j) \quad\text{subject to}\quad \sum_jg_j=\frac1T,\quad g_j\le b_j.\]This solves the one-row linear program for fixed distances. It can consume a destination needed more by a later row, so it need not solve the global problem.
For \(T=2,D=3\), source mass is \(1/2\) per step and expert capacity \(1/3\). The first step sends \(1/3\) to a full-capacity nearest sample and \(1/6\) to the next. Remaining capacities \(0\), \(1/6\), \(1/3\) total \(1/2\), matching the remaining source mass.
For fixed-length episodes, every row transports \(1/T\), capacities remain nonnegative, and total mass is one. The coupling is feasible, hence its cost upper-bounds the optimum. Early termination or variable length requires a separate mass policy; resetting capacities is insufficient.
Geometry and rewards
Positive scales \(\sigma_k\) for expert coordinate \(k\) define standardized Euclidean distance:
\[d_M(x,y)=\left(\sum_{k=1}^{d}\frac{(x_k-y_k)^2}{\sigma_k^2}\right)^{1/2}.\]Zero-variance coordinates need a positive floor or removal. Scaling changes behavioral similarity, not just numerical stability. With greedy row cost \(c_i^g=\sum_jg_{ij}d_M(x_i,y_j)\), PWIL uses
\[r_i=\alpha\exp\left(-\frac{\beta T}{\sqrt d}\,c_i^g\right), \qquad\alpha,\beta>0.\]The factor \(T\) compensates for mass \(1/T\); \(\sqrt d\) is heuristic dimensional normalization. Here \(d\) counts state/action coordinates, not points in a continuous state space. In a one-dimensional example, transporting \(1/T\) mass over distance \(2\) gives \(c_i^g=2/T\) and \(r_i=\alpha e^{-2\beta}\). Multiplying an unweighted nearest-neighbor distance by \(T\) would lose this horizon cancellation.
The surrogate objective
Monotonicity of each reward does not imply equivalence of the sums:
\[\arg\max_\pi\mathbb E\!\left[\sum_i e^{-\lambda c_i}\right] \ne \arg\min_\pi\mathbb E\!\left[\sum_i c_i\right]\]Costs \((0,2)\) and \((1,1)\) have equal sums, yet \(1+e^{-2\lambda}>2e^{-\lambda}\) for \(\lambda>0\). The exponential favors the uneven pair. Nonlinearity, greedy coupling, discounting, finite demonstrations, and policy optimization separate practical PWIL from exact transport minimization.
Remaining capacities make the reward history-dependent as a function of \((s_i,a_i)\). Adding capacities to the state yields a Markov representation; replay normally stores previously computed rewards instead. Changing normalization or reward versions therefore requires consistent relabeling or a fixed experimental version.
Computation and evaluation
Computing \(D\) distances in \(d\) dimensions costs \(O(Dd)\) per step. Sorting costs \(O(D\log D)\); repeated nearest-positive-capacity scans can cost \(O(D^2)\) per row. Amortized bounds depend on data structures and exhausted capacities. Row mass \(1/T\), nonnegative capacities, resets, tolerances, and positive scales preserve feasibility and rewards in \((0,\alpha]\) for finite distances. Tiny comparisons with an exact linear-program solver (POT) can expose errors that rising training returns miss.
Behavioral distance and environment return measure different outcomes; neither guarantees task-relevant timing. The original paper’s continuous-control experiments and state-only, normalization, replay-initialization, and capacity ablations apply to its setup, not arbitrary robots or current methods. Comparisons with behavioral cloning and another imitation method require matched demonstrations, interactions, normalization, seeds, and environments.