Robot Policy Learning with Temporal Optimal Transport Reward

In plain words

Teaching a robot a task usually means writing a reward: a score that says how well it is doing. For many tasks that score is hard to design. A common shortcut is to show the robot a video of someone doing the task and reward it for looking similar.

One way to measure "looks similar" is optimal transport, which asks how much effort it takes to move one set of frames onto another. The catch is that this measure does not care about order. A robot that does the right moves in the wrong sequence can look just as good as one that does them in order.

This paper adds time back in. Each frame is compared together with the frames right after it, and only matches that stay close in time are allowed. The result is a more accurate score from only a couple of demonstrations, and a reinforcement learning agent can then learn from that score instead of a hand-written reward.

TL;DR

• Learns robot policies from one or two expert videos with no task reward, by relabelling the agent's trajectories with an optimal-transport (OT) reward §3.1.
• Adds two changes to the OT reward: a cost averaged over a short window of frames eq. 6, and a banded temporal mask on the transport plan eq. 7, eq. 9.
• Author-reported: average success 61.1 over 9 Meta-World tasks, against 35.9 for the ADS baseline and 20.8 for the sparse task reward tab. 1.
• Costs: the OT reward needs a fixed ResNet50 encoder and a Sinkhorn solve per trajectory per demo; 1 h (state) and 3 h (pixel) on one RTX 3090 app. B.4.
• Code exists, but the repository has no license file. Use the paper's equations and the pitfalls below, not the code.

Problem

Reinforcement learning needs a reward, and reward design is hard and often done by hand §1. Imitation from a few demonstrations avoids that, but the demonstrations must be diverse and of high quality §1. Prior work uses optimal transport between agent and expert observations to make a proxy reward from only a few demonstrations §1, §3.1. The paper argues that this OT reward is invariant to the order of frames, so it can rank a correct trajectory and a wrong one with the same frames in a different order as equal §3.1.2. The OT reward for a step also depends on the whole trajectory, which makes it non-stationary during training §3.1.2. The method assumes that the agent and the expert move at a similar speed §3.2.3, app. A.4.

Method

The base reward compares an agent trajectory with an expert trajectory by optimal transport. A transport plan \mu^* is found with Sinkhorn iterations §3.1, and each agent step receives the negative cost it carries under that plan eq. 5. TemporalOT changes both stages of this pipeline §3.2.1:

1. Context cost (stage 1). The pairwise cosine cost is replaced by the average cosine cost over k_c consecutive frame pairs, shifted together in the agent and the expert sequence eq. 6 §3.2.2.
2. Temporal mask (stage 2). The transport plan is restricted to a band of half-width k_m around the diagonal eq. 9 §3.2.3. The masked problem eq. 7 is solved through its Lagrangian eq. 8, which gives a Sinkhorn-type scaling update app. C.
3. Reward and demo choice. The reward of each step is computed with the context cost and the masked plan. If several demonstrations are used, the one with the largest summed reward is kept alg. 1.
4. Policy learning. Transitions are stored with the OT reward, and the agent is updated with the TD loss eq. 1 alg. 1. The pixel setup follows the DrQ-v2 model of ADS §4.6, and the parameter table is in app. B.3.

The equations the implementation depends on:

r^{OT}_i = - \sum_{j=1}^{T} c(o_i, o^E_j)\, \mu^*(i, j) eq. 5

\hat c(o_i, o^E_j) = \frac{1}{k_c} \sum_{h=0}^{k_c - 1} \left( 1 - \frac{\langle f(o_{i+h}), f(o^E_{j+h}) \rangle}{\Vert f(o_{i+h}) \Vert\, \Vert f(o^E_{j+h}) \Vert} \right) eq. 6

\mu^* = \arg\min_{\mu} \langle M \odot \mu, \hat C \rangle_F - \epsilon\, \mathcal H(M \odot \mu), \quad \text{s.t.}\ \mu \mathbf 1 = \mu^\top \mathbf 1 = \mathbf s eq. 7

M(i, j) = \begin{cases} 1, & j \in [i - k_m, i + k_m] \\ 0, & \text{otherwise} \end{cases} eq. 9

Here f is a fixed visual encoder, \mathbf s = [1/T, \dots, 1/T] §3.2.3, and \epsilon is the entropy weight. The paper does not give a value for \epsilon §3.2.3.

The pseudocode below is our own restatement of eq. 5-9 and app. C, with Algorithm 1 alg. 1 as the outer loop. Step 2 is the context cost eq. 6; step 3 is the band mask eq. 9.

# Context cost between agent and expert frames  [eq. 6] [§3.2.2]
# Input:  traj_frames (T, ...), demo_frames (T, ...), k_c
# Output: cost (T, T) with cost[i, j] in [0, 2]
def context_cost(traj_frames, demo_frames, k_c):
    # Embed both sequences with the fixed encoder f  [§4.1] [alg. 1]
    agent_emb = f(traj_frames)                 # (T, d)
    expert_emb = f(demo_frames)                # (T, d)

    # Average the cosine cost over k_c shifted frame pairs  [eq. 6]
    cost = zeros(T, T)
    for h in range(k_c):
        # Shift both sequences by h; indices clip at the last frame
        # [code: models/temporalot.py#L287-L295]
        agent_h = shift(agent_emb, h)          # (T, d)
        expert_h = shift(expert_emb, h)        # (T, d)
        cost = cost + cosine_cost(agent_h, expert_h)
    return cost / k_c

# Band mask, 1 where |i - j| <= k_m  [eq. 9] [§3.2.3]
# Input:  T, k_m
# Output: mask (T, T)
def band_mask(T, k_m):
    mask = zeros(T, T)
    for i in range(T):
        for j in range(T):
            if abs(i - j) <= k_m:
                mask[i, j] = 1
    return mask

# Masked entropic OT with uniform marginals  [eq. 7] [app. C]
# Input:  cost (T, T), mask (T, T), eps, n_iter
# Output: masked plan (T, T), i.e. M * mu
def masked_plan(cost, mask, eps, n_iter):
    # Gibbs kernel, zero outside the band  [app. C]
    kernel = mask * exp(-cost / eps)           # (T, T)
    # Uniform row and column marginals  [eq. 7]
    marginal = ones(T) / T                     # (T,)
    scale_u = ones(T)
    scale_v = ones(T)
    for _ in range(n_iter):
        # Alternate row and column rescaling  [app. C]
        scale_u = marginal / (kernel @ scale_v)
        scale_v = marginal / (kernel.T @ scale_u)
    return diag(scale_u) @ kernel @ diag(scale_v)

# Reward for one agent trajectory against one demo  [eq. 5] [alg. 1]
def temporal_ot_reward(traj_frames, demo_frames, k_c, k_m, eps, n_iter):
    cost = context_cost(traj_frames, demo_frames, k_c)
    mask = band_mask(T, k_m)
    plan = masked_plan(cost, mask, eps, n_iter)
    # r_i = - sum_j cost[i, j] * plan[i, j]  [eq. 5]
    return -sum(cost * plan, axis=1)           # (T,)

# Keep the demo whose summed reward is largest  [alg. 1]
def trajectory_reward(traj_frames, demos, k_c, k_m, eps, n_iter):
    rewards = [temporal_ot_reward(traj_frames, d, k_c, k_m, eps, n_iter)
               for d in demos]
    best = argmax([sum(r) for r in rewards])
    return rewards[best]

# Outer loop, following Algorithm 1  [alg. 1]
# Input:  demos, N trajectories, replay buffer D
def train(agent, demos, N, replay):
    for _ in range(N):
        # Roll out the current policy to get one trajectory  [alg. 1]
        traj_frames, transitions = rollout(agent.policy)
        # Relabel each step with the OT reward, not the task reward  [alg. 1]
        rewards = trajectory_reward(traj_frames, demos, k_c, k_m, eps, n_iter)
        replay.add(transitions, rewards)
        # TD update from replay  [eq. 1] [alg. 1]
        agent.update(replay.sample())

Hyperparameters

Values come from the paper unless the "Where" column says code. Where the paper gives no value, the row says "not stated", and the code value is listed in Notes with its citation. The code's own values are used by its training script (main.py), not by its library defaults.

Results

Author-reported. Setup: Meta-World, 9 tasks, 2 expert demonstrations, 5 random seeds, success counted when the task is solved at the last step, goal position changed every episode §4.1. Baselines: TaskReward (DrQ-v2 with the sparse task reward), BC, GAIfO, OT with discount 0.99 and 0.9, and ADS §4.2. The paper reports mean and standard deviation over seeds §4.1. Table 1 in the paper prints the standard deviation in parentheses; the caption does not say so. The paper's bold marks are not reproduced here.

Table 1: success rate, Meta-World (mean (std)) tab. 1

TemporalOT beats ADS on 8 of 9 tasks; Push is the exception §4.3. The TaskReward baseline does well only on Door-lock and Window-open, because it rarely collects a first success §4.3. BC and GAIfO overfit with two demos §4.3.

Table 2: pretraining with action-inclusive demos, Door-open, success rate (columns as printed, environment steps) tab. 2 app. A.3

Table 3: demonstrations sampled faster, success rate tab. 3 app. A.4

Other results, described only in the figures and text:

• Ablations: removing the context cost or the mask lowers performance. The mask matters more on Door-open and the context cost more on Stick-push §4.4 fig. 4.
• Key parameters: a medium context length and a medium mask width work best. Large k_c adds noise; small k_m sees too few frames. More demos help §4.5 fig. 5.
• Pixel observations: TemporalOT beats ADS, often converging faster and sometimes reaching a higher final rate §4.6 fig. 6.
• Failure cases: on Hand-insert, the brown block blends into the table and the agent copies only the arm motion §4.7 fig. 7. More failure cases on Door-open, Door-unlock, Push and Stick-push are shown in app. A.2 fig. 9.
• Dynamic mask: a variant whose window follows the lowest-cost frame eq. 10 is slightly better, but the paper keeps the fixed band for simplicity app. A.6 fig. 10.
• Encoders: ResNet50 and ResNet152 perform similarly; ResNet18 is weaker app. A.5 fig. 10.

Paper vs code

Official code: github.com/fuyw/TemporalOT at 58e59c19b3afd11986e0e734da2efa0361956c9e. Linked from the abstract. The repository has no LICENSE file, so no code may be copied; the paths below are for reference only.

• Demonstrations. The paper calls them expert video demonstrations and does not say how they were made §4.1. The code generates them with MetaWorld scripted policies and drops failed episodes [code: collect_expert_traj.py#L132-L133] collect_expert_traj.py. Only two demo files are committed, for Basketball and Door-open [code: data/expert_demos/]. The script hard-codes Basketball as the environment collect_expert_traj.py, so the other 7 tasks need the script to be changed.
• Demo frames. The saved demos keep every second frame, matching the action repeat of 2 collect_expert_traj.py. An episode of 175 steps therefore has 88 frames, and 125 steps have 63 main.py. The paper gives its lengths in environment steps app. B.1 and does not say which length the mask uses.
• Encoder for the OT reward. The paper uses ResNet50 features "for each pixel observation" §4.1. The code always runs the ResNet on the 224 x 224 render, even when the agent sees state features main.py main.py main.py. So state-based runs also use image-based OT costs.
• Context windows. Paper eq. 6 leaves the window undefined at the end of a sequence. The code clips the index at the last frame [code: models/temporalot.py#L287-L295].
• Reward scale. The paper uses the reward as written in eq. 5 eq. 5. The code multiplies it by sinkhorn_rew_scale, set on the first episode to 10 / |sum of rewards| main.py main.py. The paper does not mention this scale.
• TD target. The paper's loss is one-step with a max over next actions [eq. 1] §2.1. The code uses an actor-critic update with a sampled next action, the minimum of two critics, and 3-step returns [code: models/temporalot.py#L192-L200] main.py [code: utils/buffer_utils.py#L181-L193].
• Update schedule. Algorithm 1 updates once per trajectory alg. 1. The code updates on every environment step after 6,000 agent steps and skips every other call main.py temporalot.py.
• Exploration. Not in the paper. The code takes random actions for the first 2,000 agent steps and then adds Gaussian noise that decays from 0.4 to 0 main.py main.py.
• Actor and critic. The paper gives the widths in a table tab. 4. The code has two hidden layers of 1024 plus an output [code: models/temporalot.py#L46-L50] temporalot.py.
• Dynamic mask. The dynamic mask of eq. 10 is not in the repository.
• Library defaults. TemporalOTAgent defaults to mask_k = 2 and env_horizon = 88 temporalot.py [code: models/temporalot.py#L109]. The training script passes mask_k = 10 and the episode-dependent horizon, so use the script values main.py.
• Epsilon and iterations. The paper gives no values for \epsilon or the number of Sinkhorn iterations. The code uses 0.01 and 100 [code: main.py#L37-L38].
• Table inconsistency. Table 3 gives Button-press at 1x as 89.0 (3.8) [tab. 3], while Table 1 gives TemporalOT 92.4 (3.6) and ADS 89.0 (3.8) tab. 1. The 1x figure matches the ADS entry, not TemporalOT.

No discrepancies were found in the masked Sinkhorn update itself: the code's band matches eq. 9, the uniform marginals match §3.2.3, and the reward sum matches eq. 5 ot_utils.py ot_utils.py temporalot.py.

Pitfalls

• Numerical underflow. With \epsilon = 0.01 and costs up to about 2, the naive kernel \exp(-C/\epsilon) is about e^{-200}. That underflows in float32 (our calculation). Compute the scaling in log space, as the code does with logsumexp and a large negative offset for masked entries [code: utils/ot_utils.py#L8-L28].
• Sinkhorn convergence. The code caps iterations at 100 and checks the marginal error only every 10 steps ot_utils.py. The plan's marginals may therefore not be exact. Log the marginal error if you reproduce this.
• Scale dependence. The reward scale depends on the first episode [code: main.py#L173-L176]. Changing the first episode changes the TD targets.
• Unstated constants. Exploration, start of updates, critic target noise, reward scale and the number of random steps all come from the code, not the paper [see Hyperparameters].
• Dependence on the encoder. The paper says the method depends on the pretrained encoder, and attributes a Hand-insert failure to a brown block that is close in colour to the table §6 §4.7.
• Cost grows with demos. The OT is solved once per trajectory per demo, so cost scales with N_E §6. The feature vector is about 100,000-dimensional (2048 x 7 x 7) resnet.py, and the paper does not say whether features are pooled.
• Speed assumption. The band assumes similar speed. The paper shows a performance drop when the expert moves 3x or 4x faster app. A.4 tab. 3.
• Scripted demos. The code's demos come from scripted policies, not human video. Results with other demos are not reported.
• Single-run variance. The paper reports 5 seeds, but some standard deviations are large (e.g. Door-open 78.4 (12.4), ADS Basketball 42.2 (44.5)). Compare with the std, not just the mean.

Lineage

• Builds on: the OT reward for imitation from a few demonstrations, cited in the paper as OTR §3.1 and the ADS method (Liu et al., ICLR 2024, "Imitation Learning from Observation with Automatic Discount Scheduling") §4.2. Neither has a card here.
• Adds: order-aware context costs and a temporal mask to the OT reward §3.2.
• Compared against: TaskReward, BC, GAIfO (arXiv:1807.06158 as listed in the paper's bibliography) §4.2, OT0.99, OT0.9 §4.2, and ADS §4.2.
• Superseded by: not stated.