Scorer as Optimization Target

[mgagvani]

Rationale: We want to make the scorer (1) more robust to out-of-distribution trajectory proposals, and (2) align the gradient of the scorer such that it can directly be optimized against, or arbitrarily queried to enable methods such as MPPI.

There are 3 methods I believe could work for this. 1 and 2 are more so addressing point (1), and 3 addresses (2).

1. Random Data Augmentation

• For each train sample, we run forwards+backwards on the entire module once. Then we generate $K$ random perturbations each proposed trajectory, resulting in $N \times K$ augmented samples. For each, we do a forward+backward on the scorer only, leaving the rest of the model frozen.

2. Adversarial Data

• We use the Fast Gradient Sign Method or Projected Gradient Descent to generate an "adversarial" train sample for each given train sample. FGSM will take the gradient of the scorer output score with respect to the input trajectory. It then runs gradient ascent on that sample, producing a sample which the scorer thinks has a very high score (e.g. bad), when it is very close to the original, since $||x-x'||<\epsilon$ (the adversarial $x'$ is within a limited distance to original $x$). This should (theoretically?) be a more principled fix than random data augmentation, and it sounds more unique.

3. Sobolev Training (https://arxiv.org/abs/1706.04859)

• Sobolev Training is a method for incorporating the gradients of the function we are attempting to learn into the training objective. To understand why this is necessary, see the plot:

If we take finite differences of the scorer output over a grid of possible trajectories all perturbed by $(\Delta x_i, \Delta y_j)$, where $i,j={0,2,3,...50}$, we observe the top-left plot. On the other hand, the true ADEs on the top right have their gradient going in the opposite direction. This causes a gap between orange "score min", which is the trajectory obtained by performing gradient descent on scorer ouptuts (within 1x1 meters from the original) and blue "ade min", which is the best trajectory given we can only stay within 1x1 meters of the original proposal.

I believe a solution for this is Sobolev Training, since we have access to the $n$-th derivatives of the function we are trying to learn.

• Random Data Augmentation
• Adversarial Data
• Sobolev Training