Inner Product Regularization
for Imitation Learning

Benjamin A. Christie1, Tianyu Qiu2, Amy H. Choi2,3, David Fridovich-Keil2, Dylan P. Losey1
1Virginia Tech, 2The University of Texas at Austin, 3KAIST
Under Review
Paper arXiv Code

The first-order response to state perturbations should be locally consistent across the policy.

A robot closing a coffee machine. Under behavior cloning the policy reacts to moving the coffee machine and to moving a banana; with InneR it reacts only to the coffee machine.

Abstract

Robots that learn from demonstrations should extrapolate their policy despite state perturbations such as sensor noise, scene variation, and compounding errors. Prior works use regularizers to limit the policy's sensitivity to any perturbation at each individual state. By contrast, our insight is that nearby states should respond to the same perturbation consistently: if moving an irrelevant object does not change the policy at one state, it should not at a nearby state. We therefore develop inner product regularization (InneR). We show that the inner product on action space induces an inner product on state space through the policy Jacobian (the Fisher information for stochastic policies), whose null space contains the perturbations with no first-order effect on the policy. Regularizing each inner product alone only shrinks the policy's sensitivity. Instead, InneR clusters states with aligned inner products, averages each cluster, and minimizes the smallest eigenvalues of that average so that nearby states share a common null space. InneR thus enforces consistent responses without limiting policy expressiveness, and can be appended to existing imitation learning pipelines for deterministic or stochastic policies. Our experiments compare InneR to behavior cloning and norm regularization methods. In simulated experiments, InneR makes learned policies more robust to irrelevant state dimensions and to shifts between training and test distributions, whereas the norm penalties often perform worse than behavior cloning. On a real robot arm performing three manipulation tasks, InneR raises the average success rate of a diffusion policy from 18% to 57% without additional data.

Method

Let \(\Delta\) be a small perturbation of the state \(s\). To first order the policy's action changes by \(J(s)\Delta\), where \(J(s) \in \mathbb{R}^{m \times n}\) is the Jacobian of the policy with respect to the state. If we measure distances between actions with an inner product \(M\), then

\[ \|\pi_\theta(s+\Delta) - \pi_\theta(s)\|_M \approx \|\Delta\|_{A(s)}, \qquad A(s) = J(s)^T M J(s). \]

\(A(s)\) is an inner product matrix on state space. Perturbations in its null space have no first-order effect on the action. For a stochastic policy the same role is played by the Fisher information with respect to the state,

\[ F(s) = \mathbb{E}_{a \sim \pi_\theta(\cdot \mid s)}\big[ g\, g^T \big], \qquad g = \nabla_s \log \pi_\theta(a \mid s), \]

which we estimate for diffusion policies from the gradient of the evidence lower bound.

We want nearby states to ignore the same perturbations, i.e., to share a null space. The average \(\bar{A} = \tfrac{1}{k}\sum A(s)\) over a set of states has a null space equal to the intersection of the individual null spaces, so its \(n-m\) smallest eigenvalues measure how far those states are from sharing one.

The directions that matter can change over a task, so InneR does not average over the whole dataset. For every state \(s_i\) in a training batch we form a region \(\mathcal{B}_i\) from the \(k\) states whose inner products are best aligned with \(A(s_i)\), using the distance \(\|V_iV_i^T - V_jV_j^T\|_F\) between the row spaces of their Jacobians. With \(\lambda^i_1 \geq \dots \geq \lambda^i_n\) the eigenvalues of the region's average \(\bar{A}_i\), the loss is the fraction of each average that lies in the shared null space:

\[ \mathcal{L}_\mathrm{InneR}(\theta) = \frac{1}{B} \sum_{i=1}^{B} \frac{\sum_{j=m+1}^{n} \lambda^i_j}{\operatorname{sg}\big(\sum_{j=1}^{n} \lambda^i_j\big)}, \qquad \mathcal{L}(\theta) \leftarrow \mathcal{L}(\theta) + \gamma\, \mathcal{L}_\mathrm{InneR}(\theta). \]

Frobenius and nuclear norm penalties act on each \(J(s)\) by itself. They make the policy less sensitive at that state, but they do not align the null spaces of different states.

Example

Twelve states with a two-dimensional state and a one-dimensional action; every state responds to one direction and ignores the other. The states come from two parts of a task that depend on different directions.

1. Sensitivity at each state

The black tick is the direction a state responds to, the row space of its Jacobian in \(A(s_i) = J(s_i)^T M J(s_i)\). The gray strokes follow the direction that is ignored, and the orange strokes the direction that should be ignored. After standard behavior cloning the two disagree, because the data is noisy.

2. Regions

For each state we take the \(k = 6\) states whose inner products are best aligned with it, measured by \(\|V_iV_i^T - V_jV_j^T\|_F\). The shaded states are the region of the selected state. Click another state to see its region.

3. Average

We average the inner products over the region, \(\bar{A}_i = \tfrac{1}{k}\sum_{s \in \mathcal{B}_i} A(s)\). The ellipse is this average and the bars are its eigenvalues. \(\lambda_2\) is not zero because the states in the region ignore different directions.

4. Penalty

InneR adds \(\lambda_2 / \operatorname{sg}(\lambda_1 + \lambda_2)\), averaged over the regions, to the imitation loss. Training rotates the states of each region onto a common direction, and the gray strokes line up with the orange ones.

5. Try it

Change \(\gamma\) and \(k\) below. Drag the white dot to probe another point. Drag a state to turn the direction its data gives. With \(k = 12\) there is no clustering and both parts of the task are pulled onto one direction.

Training

The same example trained from a random initialization, with the InneR term in the loss from the first step. Both policies start from the same directions and take the same number of gradient steps.

Behavior cloning fits the noisy data exactly, so its imitation loss goes to zero and its error stays. InneR ends with a higher imitation loss and a lower error. A large \(\gamma\) can pull the directions together before the imitation loss has sorted them out; try \(\gamma = 3\) with a few initializations.

Simulated Experiments

BC Frobenius Nuclear InneR-C (no clustering) InneR

Single integrator

Drone landing

Real-World Experiments

We trained diffusion policies for three tasks on a Franka arm: closing a coffee maker, stacking two cubes, and pushing a pot. An apple, an orange, and a banana act as distractors. All methods use the same demonstrations and the same architecture.

On the robot RGB-D camera fixed view Track objects SAM 3, EdgeTAM State gripper and object poses Diffusion policy DDIM Franka arm target 0.5 s ahead 10 Hz Training Demonstrations 25 to 50 per task Track objects same as on the robot Rotated copies 100 to 300 per task Imitation + γ · InneR training loss trains

Objects are tracked in the camera image and their poses, together with the gripper's, form the state. The baselines use the same pipeline with a different training loss.

Three distractors, 3× speed. Objects are tinted by the masks generated by EdgeTAM.

Success rate

The number of distractors is the number of fruit whose pose is in the policy's state. Whiskers are standard error.

Observation noise

Close with three distractors. Gaussian noise is added to every object position that the policy observes.

Sensitivity to each object

Change in the commanded gripper position when one object is displaced by 1 cm, for the three-distractor policies at states from the real-world trials. Selectivity is the mean over task objects divided by the mean over distractors.

BibTeX

@article{christie2026inner,
  title   = {Inner Product Regularization for Imitation Learning},
  author  = {Christie, Benjamin A. and Qiu, Tianyu and Choi, Amy H. and Fridovich-Keil, David and Losey, Dylan P.},
  year    = {2026},
  note    = {Under review}
}