Actions, Inverse Dynamics, and Latent Actions › Latent Action Alignment with Real Controls

Even a well-trained latent action model has no reason to place its latent actions on the same axis as the real ones. The learned $z_t$ is a knob that happens to be useful for prediction. Your downstream planner, on the other hand, has to issue a real command.

Alignment is the process of connecting a learned latent action to a real action. It sounds like a small engineering step, but it is a fundamental limitation. Unlabeled latent action discovery alone leaves a free reparametrization, and alignment must break that symmetry.

This is worth pausing on. The need for alignment is not a failure of any particular algorithm. It is a mathematical property of the problem. From the perspective of unlabeled data alone, the latent action space is defined only up to a symmetry group, and any element of the orbit is as good as any other. Breaking the symmetry requires information that is not in the unlabeled data, which is why every alignment method must import something external: a small labeled set, an assumption about statistics, a shared structure across views, or a prior on the action space.

Linear systems make the problem crisp. If the true dynamics are

$$s_{t+1} = A s_t + B a_t,$$

where:

• $s_t$: the true state at time $t$
• $a_t$: the true action at time $t$
• $A$: the true state transition matrix
• $B$: the true action matrix
• $s_{t+1}$: the true next state at time $t+1$

and a latent action model learns $s_{t+1} = A s_t + \tilde{B} z_t$, then any $z_t = M a_t$ works for any invertible matrix $M$ that keeps $\tilde{B} M = B$. In particular, the two models agree when we can write

$$\begin{aligned} A s_t + B a_t &= A s_t + \tilde{B} z_t \\ B a_t &= \tilde{B} z_t \\ z_t &= M a_t && \text{for any invertible } M \text{ satisfying } \tilde{B} M = B. \end{aligned}$$

where:

• $A$: the true state transition matrix
• $B$: the true action matrix
• $s_t$: the true state at time $t$
• $a_t$: the true action at time $t$
• $\tilde{B}$: the action matrix learned by the latent action model
• $z_t$: the learned latent action at time $t$
• $M$: an invertible matrix relating the learned and true action spaces, satisfying $\tilde{B} M = B$

The latent action is therefore identified only up to a linear reparametrization: any invertible $M$ satisfying $\tilde{B} M = B$ yields an equally valid latent action $z_t = M a_t$. On nonlinear systems the gauge freedom is even larger. If the true transition is $s_{t+1} = F(s_t, a_t)$ and the latent model learns $s_{t+1} = \tilde{F}(s_t, z_t)$, then any diffeomorphism $\psi$ of the action space gives another valid model via $\tilde{F}(s_t, \cdot) = F(s_t, \psi^{-1}(\cdot))$, with $z_t = \psi(a_t)$.

where:

• $F$: the true transition function
• $\tilde{F}$: the learned latent transition function
• $\psi$: a diffeomorphism of the action space
• $z_t = \psi(a_t)$: the latent action induced by the reparametrization

The practical consequences are concrete. Two different random seeds training the same architecture on the same data may produce latent actions that are sign-flipped, permuted, or scaled versions of each other. If the model is then used for planning, the policy that was trained on top of one version will fail on the other. If you have two robots, each with a latent action space, they may not be directly comparable even if the physical systems are pixel-identical.

For discrete latent actions, the analogue is a permutation of the codebook: code 7 might mean "jump forward" in one training run and "duck" in another. For continuous latent actions, the ambiguity is both a linear map and a smooth bijection, which is even harder to detect from data alone. Estimating the misalignment from data is possible when enough paired samples exist, but the number of samples grows with the dimension of the action space and with the richness of the diffeomorphism class you are trying to fit.

The terminology "gauge freedom" comes from physics: when two descriptions of the same physical system are related by a transformation that has no measurable consequences, the transformation is called a gauge transformation, and the space of such transformations is the gauge group. In latent action modeling, the gauge group is whatever class of reparametrizations leaves the reconstruction loss invariant. Identifying the gauge group tells you exactly what additional data or assumption is needed to fix a unique latent action space.

The most direct way to break the symmetry is to pay for a little supervision. Suppose you have a small paired dataset with both latent actions $z_t$ and real actions $a_t$. Two standard approaches follow.

• Linear probe. Fit a linear map $W$ such that $a_t \approx W z_t$. This tests whether the observed relationship is close to linear over the sampled operating range. It is cheap to fit, but it can still overfit, so alignment must be measured on held-out transitions and, ideally, under relevant distribution shifts.
• Decoder head. Add a small head $g(z_t) \to a_t$ to the latent action model and train it on the paired data while the rest of the model is frozen or slowly updated. The head absorbs the nonlinear part of the misalignment. Decoder heads are more expressive but need more paired data and can overfit if the paired set is tiny.

Both approaches work well when the paired data has the same distribution as the unlabeled data. If it does not, the alignment map will be biased. A common practical variant is to include a small amount of paired data in the original training objective, so the model learns a latent action space that is already partially aligned. This is sometimes called semi-supervised latent action learning, and it can be more sample-efficient than fitting the probe as a post-hoc step.

How many paired transitions a probe needs depends on latent dimension, conditioning, noise, model class, and distribution coverage. Learning curves on held-out trajectories are more reliable than a fixed sample-count rule. A low-dimensional linear map can require far less data than the world model, but only if the representation already exposes the relevant action structure.

When labeled pairs are scarce, distributional constraints can supplement alignment. Matching marginal statistics of latent and real actions can remove selected ambiguities, but it generally cannot recover a pointwise correspondence: many reparametrizations preserve the same distribution.

A few techniques are relevant.

• Invariance to nuisance factors. Require $z_t$ to be unchanged when the same transition is seen from a different viewpoint, background, or sensor. This forbids the encoder from encoding nuisance variables, which shrinks the gauge group.
• Cross-modal alignment. If the same action was performed, its latent action should match across two different views of the same event. This is a contrastive task with positive pairs defined by actions. Instead of labeling actions directly, you label which transitions should have the same latent action, which is often easier.
• Canonicalization. Normalize the latent action space so that $\mathbb{E}[z_t] = 0$ and $\operatorname{Cov}(z_t)$ is identity. This does not fix the rotation but it removes scale. Canonicalization is a small, cheap operation that eliminates an entire class of gauge freedom.
• Mutual information maximization. Encourage $I(z_t; a_t)$ to be high with respect to whatever partial labels exist, which is a soft version of alignment. The mutual information term is a general formulation that subsumes several specific alignment losses as special cases.

Distributional alignment is weaker than paired supervision in what it guarantees. It can leave rotations, permutations, or nonlinear symmetries unresolved and can match marginals while pairing individual actions incorrectly. Paired anchors or structural assumptions are needed when the downstream system requires a specific real-control mapping.

Discrete latent actions are a special case because alignment becomes a matching problem. If you have $K$ latent codes and $M$ real actions, and some paired data, you can solve an optimal assignment problem: match each latent code to the real action it most often co-occurs with. When $K = M$ and the mapping is bijective, this is a clean linear assignment that can be solved exactly in polynomial time by the Hungarian algorithm.

Discrete alignment is a matching problem, and the next figure shows it as a co-occurrence heatmap. The selected assignment is marked on the matrix, making the mapping from latent codes to real actions explicit.

When $K \neq M$, you have to decide whether to merge or split codes, and this is essentially a form of action abstraction.

An overcomplete latent codebook can represent several modes within one nominal real action, but simply ignoring infrequent codes may discard meaningful behavior. Utilization and transition semantics should be measured before codes are merged or pruned. When $K \neq M$, alignment is no longer a one-to-one assignment and requires an explicit many-to-one, one-to-many, or abstraction policy.

It is tempting to demand perfect alignment between latent and real actions. In practice the bar is set by the downstream task.

• Planning inside a learned simulator. An unaligned latent may suffice to optimize imagined trajectories. If the plan must be executed in the real environment, however, the chosen latent still needs a reliable map to executable controls.
• Behavior cloning. A policy can imitate latent labels without naming their semantics, but deployment still requires either a latent-to-real action decoder or an environment that accepts those same latent controls.
• Deployment on hardware. The policy has to output a valid real action. Some alignment strategy is mandatory; whether a linear probe is sufficient must be established on held-out interventions.
• Transfer across robots. Alignment has to produce compatible action spaces, which is much stronger than alignment to a single robot's actions. It requires a shared coordinate system that survives domain shift, in addition to an alignment map for each embodiment.

The required alignment depends on the deployment interface. A planner that only explores imagined rollouts has different needs from a policy that must actuate hardware. Deciding that interface early determines whether pretraining needs paired anchors or can defer the mapping to a downstream decoder.