Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Actions, Inverse Dynamics, and Latent Actions › Worked Example: Recovering Hidden Actions from Unlabeled Trajectories

Michael BrenndoerferJuly 2, 202610 min read

Part 6 of 12Learning World Models at Scale

Chapter 32Actions, Inverse Dynamics, and Latent Actions › Worked Example: Recovering Hidden Actions from Unlabeled Trajectories

Reading section 6 of 10 · Worked Example: Recovering Hidden Actions from Unlabeled Trajectories · 10 minSection 6 of 10 · 10 min

Section 6 of 10 · 10 min read

Actions, Inverse Dynamics, and Latent Actions › Worked Example: Recovering Hidden Actions from Unlabeled TrajectoriesOne section · Sections in this chapter6 / 10 ↓

Sections in this chapter10

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We can see all of the above on a small linear system that we fully control. In the noiseless equation, the true action is recoverable in closed form, so the supervised inverse model can be checked against ground truth. The learned latent model is deliberately less constrained: its result shows why reconstruction does not by itself reveal the true action coordinate.

The system is a 2D state (position and velocity) driven by a scalar force. Position integrates velocity, and velocity decays by a fixed damping factor $\gamma$ each step:

s_{t+1} = \underbrace{\begin{bmatrix} 1 & \Delta t \\ 0 & \gamma \end{bmatrix}}_{A} s_t + \underbrace{\begin{bmatrix} 0 \\ \Delta t \end{bmatrix}}_{B} a_t .

where:

$s_t = (p_t, v_t)^\top$ : the 2D state at time $t$ , with position $p_t$ and velocity $v_t$
$s_{t+1}$ : the next 2D state
$a_t \in \mathbb{R}$ : the scalar force action at time $t$
$\Delta t$ : the simulation timestep, used both for position integration and for scaling the force into a velocity change
$\gamma = 0.93$ : the per-step velocity damping factor, so that $v_{t+1} = \gamma v_t + \Delta t \, a_t$
$A = \begin{bmatrix} 1 & \Delta t \\ 0 & \gamma \end{bmatrix}$ : the state transition matrix, which advances position by velocity and damps velocity
$B = \begin{bmatrix} 0 \\ \Delta t \end{bmatrix}$ : the action matrix, which injects force only into the velocity component

The action enters only through velocity. The vector $B$ has rank one, so the effective action space is one-dimensional. In the stated deterministic dynamics, $B$ has full column rank and the feasible residual lies in its column space, so $a_t = B^{+}(s_{t+1} - A s_t)$ recovers the action exactly. The simulation below adds process noise, so recovery from each realized transition is approximate rather than exact.

We start by importing libraries and generating a synthetic dataset of 400 trajectories, each 50 steps long, with random Gaussian actions. The random actions act as a natural excitation signal. Without them, the state would not explore enough of its reachable set to identify the dynamics, a point that matters in any real data collection problem.

In[9]:

Code

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.linear_model import Ridge
from sklearn.metrics import r2_score

seed = 11
rng = np.random.default_rng(seed)
torch.manual_seed(seed)

dt = 0.1
A_true = np.array([[1.0, dt], [0.0, 0.93]])
B_true = np.array([[0.0], [dt]])

The simulation is a simple loop over the transition equation.

In[10]:

Code

def simulate(n_traj, T, rng):
    states = np.zeros((n_traj, T + 1, 2))
    actions = rng.normal(0.0, 1.0, size=(n_traj, T, 1))
    states[:, 0] = rng.normal(0.0, 0.3, size=(n_traj, 2))
    process_noise_std = 0.02
    for t in range(T):
        states[:, t + 1] = states[:, t] @ A_true.T + actions[:, t] @ B_true.T
        states[:, t + 1] += rng.normal(0.0, process_noise_std, size=(n_traj, 2))
    return states, actions


states, actions = simulate(400, 50, rng)

The next blocks flatten the trajectories into transitions, create a trajectory-level train/test split, and report the resulting counts. Splitting by trajectory rather than by transition matters: consecutive transitions are highly correlated, and a random transition split would leak the answer across the boundary.

In[11]:

Code

s_t = states[:, :-1].reshape(-1, 2)
s_next = states[:, 1:].reshape(-1, 2)
a_true = actions.reshape(-1, 1)
n_transitions = s_t.shape[0]

n_traj = states.shape[0]
T = states.shape[1] - 1
train_traj = rng.random(n_traj) < 0.8
is_train = np.repeat(train_traj, T)

Out[12]:

Console

Transitions: 20000, train: 16500, test: 3500
Trajectories: 400, train: 330, test: 70

The split is now by complete trajectories, so test transitions come from held-out trajectories and there is no leakage across time. The printed counts summarize the resulting train/test partition.

Now the supervised part: fit a linear inverse dynamics model that predicts $a_t$ from the pair $(s_t, s_{t+1})$ . This is the best-case variant of the IDM shortcut from the earlier section, played out on data where the shortcut is guaranteed to work.

In[13]:

Code

X_idm = np.concatenate([s_t, s_next], axis=1)
idm = Ridge(alpha=1e-6).fit(X_idm[is_train], a_true[is_train])
a_idm = idm.predict(X_idm[~is_train])
a_idm_test = a_idm.ravel()

idm_mse = float(np.mean((a_idm_test - a_true[~is_train].ravel()) ** 2))
idm_r2 = float(r2_score(a_true[~is_train], a_idm_test))

Out[14]:

Console

Inverse dynamics MSE (test): 0.03786
Inverse dynamics R^2 (test): 0.9609

Because the system is linear and the velocity component of $s_{t+1} - A s_t$ contains $a_t$ , the inverse model recovers the action up to the injected process noise. The small regularization ( $\alpha = 10^{-6}$ ) stabilizes the numerical fit.

Now the interesting part. We pretend we never saw $a_t$ and train a latent action model. The encoder takes $(s_t, s_{t+1})$ and outputs a 1D latent $z_t$ . The decoder takes $(s_t, z_t)$ and predicts $s_{t+1}$ . Training uses no action labels, only reconstruction of the next state. Because the encoder receives that target, reconstruction can be solved through target leakage rather than action discovery; the comparison below is designed to expose that gap.

In[15]:

Code

class LatentActionModel(nn.Module):
    def __init__(self, state_dim=2, latent_dim=1, hidden=64):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(2 * state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, latent_dim),
        )
        self.decoder = nn.Sequential(
            nn.Linear(state_dim + latent_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, state_dim),
        )

    def forward(self, s_t, s_next):
        z = self.encoder(torch.cat([s_t, s_next], dim=-1))
        s_pred = self.decoder(torch.cat([s_t, z], dim=-1))
        return z, s_pred

We now train for 500 epochs with Adam. The loss is the mean-squared error of the reconstructed next state.

In[16]:

Code

S_t = torch.tensor(s_t, dtype=torch.float32)
S_next = torch.tensor(s_next, dtype=torch.float32)
S_t_train = torch.tensor(s_t[is_train], dtype=torch.float32)
S_next_train = torch.tensor(s_next[is_train], dtype=torch.float32)

model = LatentActionModel()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
losses = []
for epoch in range(500):
    optimizer.zero_grad()
    z, s_pred = model(S_t_train, S_next_train)
    recon = ((s_pred - S_next_train) ** 2).mean()
    recon.backward()
    optimizer.step()
    losses.append(float(recon.item()))

Out[17]:

Console

Final reconstruction loss: 0.000222

The reconstruction loss should drop well below the action-induced variance. If the latent collapsed, the loss would plateau near the best state-only predictor. However, a drop below the state-only level (and even below the process-noise variance) can also come from information leakage: because the encoder sees $s_{t+1}$ , the latent can encode part of the target itself. Low reconstruction error alone is therefore not evidence of an action-aligned latent. We fit a linear probe from $z_t$ to $a_t$ using the same trajectory-level train/test split and check the held-out $R^2$ .

In[18]:

Code

with torch.no_grad():
    z_all, s_pred_all = model(S_t, S_next)
z_all = z_all.numpy()

probe = Ridge(alpha=1e-3).fit(z_all[is_train], a_true[is_train])
a_from_z = probe.predict(z_all[~is_train]).ravel()

a_true_test = a_true[~is_train].ravel()
align_r2 = float(r2_score(a_true_test, a_from_z))
align_mse = float(np.mean((a_from_z - a_true_test) ** 2))

Out[19]:

Console

Latent-to-true alignment R^2 (test): 0.0815
Latent-to-true alignment MSE (test): 0.88936

The alignment $R^2$ is low in this minimal setting. Reconstruction alone does not force the latent to be a scaled and sign-flipped version of $a_t$ ; the encoder can carry next-state information that the probe cannot map back to the action. A stronger structural assumption or paired supervision is needed to make the latent action-aligned. A scalar linear probe can correct a scale-and-sign mismatch, but the unconstrained encoder-decoder admits a much larger family of nonlinear reparametrizations and target-leaking solutions.

The alignment probe is a low-dimensional fit on top of an already structured representation. The next figure measures how test-set alignment R^2 changes as the paired training set grows, using the latent actions and true actions already computed.

In[20]:

Code

import numpy as np
from sklearn.linear_model import Ridge
from sklearn.metrics import r2_score

rng_probe = np.random.default_rng(42)
train_idx = np.where(is_train)[0]
test_idx = np.where(~is_train)[0]
n_train = train_idx.size
sample_sizes = np.array([5, 10, 20, 50, 100, 200, 500, 1000])
sample_sizes = sample_sizes[sample_sizes <= n_train]
align_r2_curve = []
for n in sample_sizes:
    chosen = rng_probe.choice(train_idx, size=n, replace=False)
    probe_n = Ridge(alpha=1e-3).fit(z_all[chosen], a_true[chosen])
    pred_n = probe_n.predict(z_all[test_idx]).ravel()
    align_r2_curve.append(float(r2_score(a_true[test_idx].ravel(), pred_n)))

Out[21]:

Visualization

Line plot of held-out alignment R-squared remaining low as the number of paired probe samples increases. — Test-set alignment R^2 as a function of the number of paired transitions used to fit the linear probe. Performance plateaus quickly in this one-dimensional latent action space, but the plateau sits at a low R^2 because the unsupervised latent is only weakly action-aligned.

The first figure shows the supervised inverse dynamics prediction on the test set.

Out[22]:

Visualization

Scatter plot of predicted versus true actions lying on the diagonal. — Inverse dynamics prediction of the true action from consecutive states. On this linear system with a single effective control channel, a simple ridge regression recovers the action accurately up to the process-noise floor.

The second figure shows the same check for the fully unsupervised latent action model, using the linear probe to compare back to the true action.

Out[23]:

Visualization

Diffuse scatter of latent-action probe predictions versus true actions, with most points far from the diagonal reference. — Linear probe from the learned latent action to the true action. In this minimal unsupervised setup the alignment is weak, illustrating that next-state reconstruction alone does not guarantee a latent that spans the true control direction.

The third figure shows the training curve of the latent action model.

Out[24]:

Visualization

Log-scale line chart of reconstruction loss decreasing over training epochs. — Reconstruction loss during training of the latent action model. The model has no access to action labels and drives the error down by passing transition information through the latent; low loss does not establish action semantics.

The results illustrate several points at once. A linear inverse dynamics model recovers the true action accurately in this controlled setting. A latent action model trained only on next-state reconstruction reaches low reconstruction error while remaining weakly correlated with the true action. A linear probe measures that weak alignment without retraining the encoder.

The same experiment becomes less identifiable with nonlinear dynamics, unobserved disturbances, or a rank-deficient action map. A second action input with a smaller transition effect may be underrepresented by a reconstruction objective even when it matters for a downstream task.