The idea in one minute. PWIL compares an agent’s state–action visits with expert demonstrations using optimal transport. It constructs a dense imitation reward by greedily transporting a small amount of agent mass to nearby expert samples. Capacity constraints encourage matching how often the expert visits each region, rather than repeatedly matching one convenient example.
Reading path. The first three sections define the mathematical problem. Then follow one transport step and the reward scaling. The final sections explain what is optimized, what is only approximated, and how to evaluate an implementation. Prerequisites are probability distributions, norms, and an RL policy.
1. From trajectories to probability measures
An episode gives state–action vectors \(x_i=(s_i,a_i)\in\mathbb R^d\), \(i=1,\ldots,T\). Expert demonstrations give vectors \(y_j\), \(j=1,\ldots,D\). The empirical measures are
\[\widehat\rho_\pi=\frac1T\sum_{i=1}^{T}\delta_{x_i},\qquad \widehat\rho_E=\frac1D\sum_{j=1}^{D}\delta_{y_j}.\]Each is a probability measure with total mass one. Matching them matches the distribution of visits, not the temporal ordering of a trajectory. Two episodes with the same multiset of state–action pairs have identical empirical measures even if their order differs. Temporal constraints must therefore come from the environment, the policy, or an enriched representation.
2. Wasserstein distance and feasible couplings
For a 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).\]\(\Pi(\mu,\nu)\) contains joint probability measures whose marginals are \(\mu\) and \(\nu\). The coupling decides how much probability mass moves from each source to each destination. For the uniform empirical measures above, it is a nonnegative matrix \(\Gamma\in\mathbb R_+^{T\times D}\) satisfying
\[\sum_{j=1}^{D}\Gamma_{ij}=\frac1T, \qquad \sum_{i=1}^{T}\Gamma_{ij}=\frac1D.\]A rectangular coupling is not generally a doubly stochastic matrix with row and column sums one. Those sums would be incompatible when \(T\ne D\). The probability marginal constraints are the essential definition.
The discrete transport 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.\]For equal one-point measures, \(W_p(\delta_x,\delta_y)=d_M(x,y)\). This is a useful unit check for both the power and normalization.
3. The imitation objective
The conceptual goal is to choose a policy that minimizes expected discrepancy between its visits and expert visits. A trajectory-based empirical objective is
\[\min_\pi\;\mathbb E_{\tau\sim\pi} \left[W_1(\widehat\rho_\pi(\tau),\widehat\rho_E)\right].\]This is not automatically identical to a Wasserstein distance between infinite-horizon occupancy measures; expectation and optimal transport do not commute freely. PWIL uses a practical episodic construction. Its main appeal is that reward construction does not require fitting an adversarial discriminator alongside the policy. Dadashi et al., PWIL.
Given an exact optimal coupling for a completed episode, define the row cost
\[c_i^*=\sum_j\Gamma^*_{ij}d_M(x_i,y_j).\]Then \(\sum_i c_i^*=W_1\). But the exact row costs depend on the whole trajectory, including future states; they are unavailable when deciding an early action online.
4. Greedy transport with residual capacity
At the start of an episode, give each expert sample residual capacity \(b_j=1/D\). At agent step \(i\), move source mass \(1/T\) into the nearest available expert samples, reducing their capacities. The step solves
\[\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.\]For fixed distances, filling destinations in increasing distance order solves this one-row linear program. It does not solve the global multi-row transport problem, because an early decision can consume a destination that a later point needs more.
With \(T=2,D=3\), each agent step supplies \(1/2\) mass and each expert sample accepts \(1/3\). If the nearest destination has full capacity, the first step sends \(1/3\) there and \(1/6\) to the next destination. Residual capacities after the step are \(0\), \(1/6\), and \(1/3\), totaling \(1/2\), exactly the remaining source mass.
Across a fixed-length episode, the resulting matrix is feasible: every row supplies \(1/T\); capacities never become negative; total transported mass is one. Hence its transport cost is an upper bound on the optimal transport cost. Early termination and variable episode length need an explicit mass policy; resetting capacities alone does not resolve that issue.
5. Standardized geometry and reward scaling
Let \(\sigma_k\) be a positive scale derived from expert coordinate \(k\). A standardized Euclidean distance is
\[d_M(x,y)=\left(\sum_{k=1}^{d}\frac{(x_k-y_k)^2}{\sigma_k^2}\right)^{1/2}.\]For zero-variance coordinates, use a documented positive floor or remove the coordinate. Otherwise the distance is undefined. Different feature scales encode different judgments of behavioral similarity; they are not merely a numerical implementation detail.
Define greedy row cost \(c_i^g=\sum_jg_{ij}d_M(x_i,y_j)\). PWIL’s positive reward form is
\[r_i=\alpha\exp\left(-\frac{\beta T}{\sqrt d}\,c_i^g\right), \qquad\alpha,\beta>0.\]The \(T\) factor compensates for the row mass \(1/T\). The \(\sqrt d\) factor is a heuristic dimensional normalization, not a universal theorem about distances. Here \(d\) is the number of state and action coordinates, not the cardinality of a continuous state space.
For a toy one-dimensional row that moves \(1/T\) mass a distance \(2\), \(c_i^g=2/T\) and \(r_i=\alpha e^{-2\beta}\). The horizon cancels. This example catches an implementation that accidentally uses an unweighted nearest-neighbor distance and then multiplies by \(T\) again.
6. Why the reward is a surrogate
Even if a reward decreases with each row cost, maximizing the sum of exponentials is not generally equivalent to minimizing the sum of row costs:
\[\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]\]in general. For example, costs \((0,2)\) and \((1,1)\) have equal sums, but \(1+e^{-2\lambda}>2e^{-\lambda}\) for \(\lambda>0\). The exponential prefers the uneven pair in that example. Policy optimization, discounting, greedy coupling, finite demonstrations, and this nonlinear transform all distinguish practical PWIL from exact global Wasserstein minimization.
The reward also depends on remaining expert capacities, so it is history-dependent when viewed only as a function of \((s_i,a_i)\). One can augment the state with residual capacities to obtain a Markov representation, but an ordinary replay buffer usually stores the previously computed reward rather than reconstructing that augmented process. Keep reward version and normalization fixed during an experiment, or relabel stored transitions deliberately.
7. Complexity and implementation choices
Computing \(D\) distances in \(d\) dimensions costs \(O(Dd)\) per step. Sorting all destinations costs \(O(D\log D)\); repeatedly scanning for the nearest positive-capacity point can cost \(O(D^2)\) in the worst case for one row. Tighter amortized bounds depend on the data structure and how many capacities are exhausted. Do not label a particular paper’s implementation bound as the cost of every greedy transport algorithm.
An implementation should preserve: row mass exactly \(1/T\), nonnegative residual capacities, episode resets, floating-point tolerance, positive feature scales, and reward values in \((0,\alpha]\) for finite distances. Checking total transported mass and comparing tiny examples to an exact linear-program solver gives stronger evidence than only observing that an RL training curve rises.
8. Evaluation and current scope
Compare behavioral distance and environment return separately. Environment return can be high even when imitation differs; a low empirical transport cost can also fail to capture task-relevant timing. Include behavioral cloning and another imitation baseline under the same demonstrations, interactions, normalization, seeds, and evaluation environment.
The original PWIL paper reports continuous-control experiments and ablations on state-only inputs, normalization, replay initialization, and expert capacity. Those are historical findings under its setup. They do not establish universal superiority over current imitation methods or deployment readiness on a different robot.
References
- PWIL, 2020 / ICLR 2021.
- POT: Python Optimal Transport: exact and regularized transport tools for small validation problems.
- Acme: RL components and the historical PWIL implementation ecosystem.