DEV Community

stepinto.dev
stepinto.dev

Posted on Edited on Originally published at stepinto.dev

Latent State Is the Fields You Never Logged

The last post ended with a challenge. Four weights fitted from logs can learn a room, but give that room a door that opens and closes on a schedule the logger never records, and the best a linear model can do is average over the door.

The log has three columns: the room temperature, the outside temperature and whether the heater is on. When the door opens the room cools faster, and every row written while it is open looks like a row from a room that is simply worse at holding heat. The model sees the effect and never the cause, so fitting the same three columns harder gets it nowhere. What the log does still have is order. A room that was cooling steadily and then drops three degrees in one step has told you something, even if no column says what.

The 102 post described how we might work around this: give the model a private field it keeps for itself, and let it decide what to put there. It also promised to explain how such a field can be learned when nobody ever sees it.

This post builds the room with the door, watches the four-weight model fail on it, gives the model one private field, and watches the field turn into the door. The maths comes at the end, because by then it describes something you have watched run. The playground has the door in it now, so everything below can be rerun in your browser.

An invisible door (according to the data)

I'm using the same simulated room from the last post, with one addition. Someone keeps opening the door, and while it is open the room leaks heat three and a half times faster.

const LEAK_CLOSED = 0.07;   // the 103 room
const LEAK_OPEN = 0.25;
let doorOpen = false;       // the room's secret; never written to the log

function someone() {
  // opens the door about once every 25 steps, leaves it open about 8
  if (doorOpen) { if (rand() < 1 / 8) doorOpen = false; }
  else if (rand() < 1 / 25) doorOpen = true;
}

function realRoom(state: State, action: Action): State {
  const heaterOn = heaterAfter(state, action);
  const leak = (state.outsideTemp - state.roomTemp) * (doorOpen ? LEAK_OPEN : LEAK_CLOSED);
  const heat = heaterOn ? REAL_HEAT : 0;
  return { roomTemp: state.roomTemp + leak + heat + gauss() * NOISE, outsideTemp: state.outsideTemp, heaterOn };
}
Enter fullscreen mode Exit fullscreen mode

The rule from the last post still holds: nothing below this line looks at doorOpen. The logger calls someone() before every step and writes down the same three fields it always did, so a row from the new log is indistinguishable from a row from the old one. Eighty episodes of sixty steps gives 4,800 transitions, and the door is open for 22 percent of them. Nothing in the log says which.

If you run devices in the field, this is every column your telemetry does not have. The window someone left open, the second heater in the room that is not on your network, the pump that was switched to manual for an afternoon. The dynamics change and the log does not say why.

What four weights do with a door

First I refitted the four-weight model on the new log, with the same features, the same twenty-line training loop and ten thousand epochs.

weight sealed room (103) with the door truth, door closed
room -0.074 -0.037 -0.07
outside 0.069 0.040 0.07
heater 1.19 0.97 1.2
bias 0.077 -0.41 0

I expected the loss to get worse, and it flattens at 0.29 against 0.01 in the sealed room. I didn't expect the shape of the damage. The model has not averaged over the door, which would have given a room weight somewhere between the closed and open leak rates. It has halved the room weight and gone looking for the lost heat in the bias.

The reason is that the door is not independent of the things the model can see. An open door drags the room down, so in the log a cold room and a big drop arrive together, and the only lever the model has for "cold room" is the room weight. It pushes that weight towards zero to account for the drops, then pulls the bias negative to make the sums balance. A hidden variable that moves with the visible ones bends the fit, which is rather worse than adding noise to it.

You can see the bend when the model is asked to forecast. Later in the post there is a room at minus four with the heater on, where the door opens at step 10. Ask this model at step 13, with the room at 7.8 degrees and falling fast, what the next five steps look like, and it says the room warms to 8.3. The real room goes to 2.7.

One private field

The 102 post sketched the shape of the fix as a class with a field nobody outside can read:

class WorldModel {
  #state: Float32Array;  // s_t lives here, and nowhere in your logs
  observe(o: Observation) { /* fold the new reading into #state */ }
  act(a: Action): Distribution<Observation> { /* predict, given #state and a */ }
}
Enter fullscreen mode Exit fullscreen mode

Here is the smallest version of that I could make work. The private field is one number. act is the four-weight model from 103, plus a second weighted sum that the field switches on. observe folds what just happened into the field, using a weighted sum of its own, squashed through tanh so the field stays between minus one and one.

type Weights = {
  // act: change in room temperature = base + field * extra
  room: number; outside: number; heater: number; bias: number;
  fieldRoom: number; fieldOutside: number; fieldBias: number;
  // observe: next field = tanh(a weighted sum of the field and what just happened)
  uField: number; uRoom: number; uOutside: number; uHeater: number; uChange: number; uBias: number;
};

class WorldModel {
  #field = 0;              // s_t lives here, and nowhere in the logs
  w: Weights;
  constructor(w: Weights) { this.w = w; }

  reset() { this.#field = 0; }

  act(state: State, action: Action): State {
    const f = inputs(state, action), w = this.w;
    const base = w.room * f.room + w.outside * f.outside + w.heater * f.heater + w.bias;
    const extra = w.fieldRoom * f.room + w.fieldOutside * f.outside + w.fieldBias;
    return { roomTemp: state.roomTemp + base + this.#field * extra, outsideTemp: state.outsideTemp, heaterOn: f.heaterOn };
  }

  observe(state: State, action: Action, nextState: State) {
    const f = inputs(state, action), w = this.w;
    const change = (nextState.roomTemp - state.roomTemp) / 5;
    this.#field = Math.tanh(
      w.uField * this.#field + w.uRoom * f.room + w.uOutside * f.outside +
      w.uHeater * f.heater + w.uChange * change + w.uBias);
  }
}
Enter fullscreen mode Exit fullscreen mode

Thirteen numbers, against four. Seven of them say what the field does to the prediction; six say what the model should write into the field after each step. Nothing in there mentions a door. The field is allowed to be anything between minus one and one, and the model gets to decide what it means.

Two small things changed alongside. inputs is features with the temperatures centred and scaled (room temperature minus 17, outside minus 7, both divided by five), and the change the model observes is divided by five too. In the last post I left the inputs raw and paid for it with ten thousand epochs down a long thin valley. This time the loop below is more expensive per epoch, so I took the shortcut the 103 post mentioned, and it converges in three thousand.

The other change is to the observe signature. The 102 sketch took an observation; this one takes the whole transition, because the thing the field most needs to know is how the room moved compared with what the model expected, and that needs both ends of the step.

Training through the field

The 102 post promised to explain how this gets trained. The field is never in the log, so there is no column to compare it against. The thirteen weights have to be judged by the only thing that is in the log, the next temperature, and six of those weights only affect the next temperature indirectly, by changing what gets written into the field, which changes the prediction a step later, and the step after that.

So the gradient has to flow through the field. In 103 each row of the log was its own small problem, and the loop could visit rows in any order. Now a row's prediction depends on the field, the field depends on the previous row, and so on back to the start of the episode. The training loop walks each episode forwards, remembering what the field was at every step, and then walks it backwards, asking at each step how much that step's field mattered to everything that came after it.

function train(log: Transition[][], epochs: number, lr: number) {
  const w = {} as Weights;
  for (const k of keys) w[k] = k.startsWith("u") ? (rand() - 0.5) * 0.2 : 0;
  const losses: number[] = [];
  const n = log.flat().length;
  for (let epoch = 0; epoch < epochs; epoch++) {
    const g = {} as Weights;
    for (const k of keys) g[k] = 0;
    let loss = 0;
    for (const episode of log) {
      // forwards through the episode, remembering what the field was at every step
      const steps = [];
      let field = 0;
      for (const t of episode) {
        const f = inputs(t.state, t.action);
        const target = t.nextState.roomTemp - t.state.roomTemp;
        const base = w.room * f.room + w.outside * f.outside + w.heater * f.heater + w.bias;
        const extra = w.fieldRoom * f.room + w.fieldOutside * f.outside + w.fieldBias;
        const err = base + field * extra - target;
        loss += err * err;
        const change = target / 5;
        const next = Math.tanh(w.uField * field + w.uRoom * f.room + w.uOutside * f.outside + w.uHeater * f.heater + w.uChange * change + w.uBias);
        steps.push({ f, change, field, extra, err, next });
        field = next;
      }
      // backwards through the episode: how much did each step's field matter to everything after it?
      let dField = 0;   // d loss / d (field after this step)
      for (let i = steps.length - 1; i >= 0; i--) {
        const s = steps[i];
        const dPre = dField * (1 - s.next * s.next);           // through the tanh
        g.uField += dPre * s.field; g.uRoom += dPre * s.f.room; g.uOutside += dPre * s.f.outside;
        g.uHeater += dPre * s.f.heater; g.uChange += dPre * s.change; g.uBias += dPre;
        g.room += s.err * s.f.room; g.outside += s.err * s.f.outside; g.heater += s.err * s.f.heater; g.bias += s.err;
        g.fieldRoom += s.err * s.field * s.f.room; g.fieldOutside += s.err * s.field * s.f.outside; g.fieldBias += s.err * s.field;
        dField = s.err * s.extra + dPre * w.uField;            // on to the step before
      }
    }
    for (const k of keys) w[k] -= lr * g[k] / n;
    losses.push(loss / n);
  }
  return { w, losses };
}

const { w, losses } = train(log, 3000, 0.3);
Enter fullscreen mode Exit fullscreen mode

The forwards half is act and observe inlined. The backwards half is the chain rule, and two lines of it carry the whole idea. dPre is how much the loss would change if the field written at this step were nudged, pulled back through the tanh. The last line passes that on to the step before, because the step before chose the field this step started with. Everything else is the err * input pattern from 103, applied to whichever weight the line is about.

This is backpropagation through time, and it is the reason recurrent models are trained on episodes rather than on rows. Three thousand epochs over 4,800 rows takes about half a second in Node 23. The six u weights start at small random values rather than zero, because with every u at zero the field never moves, nothing depends on it, and the gradient for those weights is exactly zero forever. The seven prediction weights start at zero as before.

Mean squared error while training the latent model, three thousand epochs

Mean squared error per epoch for the latent model, with three reference lines. The four-weight model stopped at 0.29. A model that could see the door would reach 0.01, the level of the draughts. The same model told the door one step late gets 0.17, and the latent model settles just under it.

Download the data (CSV)

The two dashed lines show what the model is being asked to do. I fitted a seven-weight model that is told the door state, which is cheating, to find out what the best possible fit looks like. Scored normally it reaches the noise floor. Then I gave the same model the door state one step late, so it only learns the door has moved after seeing the step on which it moved, and it gets 0.174. The step the door moves on is a surprise by definition, and the 310 steps where it moved account for 0.165 of that on their own. The latent model gets 0.172.

That second line is a reference rather than a floor, and the latent model sitting under it is the clue. The late-door model insists the door is wherever it was a step ago. A model that knows the door a step late and also knows how often it moves (about one step in eight while it is open, one in twenty-five while it is shut) can hedge, adding a little leak to every closed step in case the door has just opened, and that scores 0.163. Refitting the seven weights on the late door state lands in the same place. So the latent model is about a hundredth above what noticing the door can reach with this shape of model. None of these numbers is a true lower bound, either: all of them are scored on the 4,800 rows the models were fitted to, and a model with a better shape could go lower.

Here is what it learned, with the temperature weights divided by five so they read per degree.

weight learned weight learned
room -0.111 uField 0.007
outside 0.078 uRoom -0.049
heater 1.21 uOutside 0.031
bias -0.92 uHeater 0.50
fieldRoom 0.010 uChange -2.19
fieldOutside 0.015 uBias -0.39
fieldBias -1.83

Read the right-hand column as the recipe for the field. The weight on uChange is by far the largest: the field is mostly a function of how much the room moved last step, with the heater and the temperatures subtracted off so that a drop the base model already expected does not count. The weight on uField, the field's memory of itself, is almost nothing. For this room the model decided it did not need memory, because one step of an open door is loud enough to hear on its own. That changes when the door gets quieter, further down.

The left-hand column is less tidy than in 103. The room and outside weights are no longer equal and opposite, and the field's effect is almost all in fieldBias. The model is not obliged to factor the room into "closed door physics plus a door term" the way I would, and the gradient found a different split that fits the log just as well. A private field has no label on it. It holds whatever the model found useful, and you judge it by the predictions, as the 102 post said you would have to.

The field finds the door

The test is the same unseen room as in the last post: 22 degrees inside, minus four outside, colder than anything in the log. The heater is on throughout. The door opens at step 10 and closes at step 20, and the models are not told.

Two things are different about how the models are run. First, each one sees the real reading at every step and predicts the next one, the way a thermostat would, rather than running thirty steps blind from the start. With a door it cannot see, no model could run thirty steps blind, and the point of the field is that the model notices. Second, at three chosen steps, each model is asked to run five steps blind from where it is, which is the "what if" that world models are for.

Thirty steps with the heater on at minus four outside, the door open from step 10 to 20

Top: the real room, with five-step forecasts from both models branching off at steps 5, 13 and 23. Bottom: the private field after each step, with the door state it never saw dashed over it. The field is near zero for ten steps, 0.91 the step after the door opens, and 0.12 the step after it closes.

Download the data (CSV)

The strip at the bottom is the result I was after. The field sits a little below zero while the door is closed, jumps to 0.9 on the first step after it opens, stays high while it is open, and drops back the step after it closes. The model was never told there was a door, and the field is a door anyway, because a door was the thing in the room that the log could not explain and a single number could.

Step by step at the two moments that matter:

step door real room four-weight, one step ahead latent, one step ahead field after
10 closed 17.19 17.16 16.97 -0.15
11 open 13.14 16.96 16.90 0.91
12 open 10.10 13.06 11.16 0.86
13 open 7.76 10.13 8.52 0.81
20 open 1.94 2.69 2.27 0.59
21 closed 2.83 2.27 1.77 0.12
22 closed 3.55 3.13 3.56 0.15

Step 11 is wrong for both models and has to be. Nobody could have predicted the door. From step 12 the latent model is within a degree and closing, while the four-weight model is three degrees adrift and stays there. Step 21 is the same story in reverse. Over the thirty steps the one-step error is 1.16 degrees for the four-weight model and 0.78 for the latent one; leave out the two steps the door moved on and it is 0.95 against 0.32.

At step 13 the latent model has had two steps to notice the door, and its field reads 0.81. Asked for five steps blind, it says the room will fall from 7.8 to 2.7, and the room falls to 2.7. The four-weight model, from the same reading, says 8.3. It has no way to represent "the room is currently leaking faster than usual", so it cannot condition on it, and its bent physics sends the forecast the wrong way.

forecast from real, five steps on four-weight latent
step 5, door closed 17.2 17.7 17.1
step 13, door open 2.7 8.3 2.7
step 23, door just closed 6.8 5.3 6.1

The strip also shows the field drifting down from 0.91 to about 0.55 over the ten open steps. It is reading the size of the drop, and the drop shrinks as the room gets closer to the temperature an open room would settle at. The model has no memory to hold the field up because, as the weights table said, it did not learn any. Make the door quieter and that changes. Rerunning the whole script with the open-door leak at 0.15 instead of 0.25 brings the memory weight from 0.007 to 0.08, and at 0.12 it is 0.26: when one step is no longer enough to tell, the model learns to add up several. At 0.10, where the door only adds a few hundredths to the leak, the field stops tracking the door at all, and the latent model's loss sits above the one-step-late line. There is a signal it cannot find in 4,800 rows, and the answer, as with most problems in this series, is more log.

Now the notation

Everything above fits in two lines. The private field is (h), the readings the logger can see are (o), and the action is (a).

$$h_t = f(h_{t-1},\ o_{t-1},\ a_{t-1},\ o_t)$$

$$\hat{o}_{t+1} = g(h_t,\ o_t,\ a_t)$$

The first line is observe: the new field is a function of the old field and the step that just happened. The second is act: the predicted next reading is a function of the field, the current reading and the action. In this post (f) is a tanh of a weighted sum and (g) is two weighted sums, and training them was gradient descent with the chain rule run backwards through the first line.

Put the field and the readings together and call the pair (s_t), and the second line is the world model equation from the 102 post, (P(s_{t+1} \mid s_t, a_t)), with the probability dropped because the model predicts a single number. The state the equation conditions on was always meant to include the fields you never logged. That is why it is called latent: the field is hidden from the logger, but the maths had a place for it all along.

How DreamerV3 does the same job

The part of this post that carries over most directly to the research systems is the one it spent longest on. Nobody tells DreamerV3 what to keep in its recurrent state either. The state has to hold whatever a single observation cannot show, such as how fast something is moving or where an object went when it left the frame, and it learns to hold those things for the same reason the field learned the door: the predictions were wrong without them, and the gradient, walking back through each sequence, found which numbers would fix that.

The architecture carries over less neatly. DreamerV3 calls its world model a recurrent state-space model, and it has more parts than my two methods:

  • a sequence model, a GRU, that carries a recurrent state (h_t) forward from (h_{t-1}), the last action and the last latent (z_{t-1})
  • an encoder that reads the new observation alongside (h_t) and produces (z_t), a stochastic latent made of 32 categorical variables with 32 classes each
  • a dynamics predictor that guesses (z_t) from (h_t) alone, before the observation arrives
  • a decoder, with heads for reward and for whether the episode carries on, that reads (h_t) and (z_t) together

My observe does the work of the sequence model and the encoder in one line, and my act is closest to the decoder. Three of the differences matter, and each of them shows up somewhere in this post.

The first is where the readings go. My act is handed the current reading, so the field only has to carry what the reading leaves out, and one number covers the door. DreamerV3's decoder has nothing but the state to work from, so the state has to carry everything, the equivalent of the temperature included. That is why it is hundreds or thousands of numbers wide, and it is what lets the model imagine many steps ahead with no observations at all. My five-step forecasts get away with feeding the predicted temperature back in as the next reading, which works only because the temperature is the whole of what the logger sees.

The second is the guess made before the reading. DreamerV3 keeps two versions of (z_t) at every step: the dynamics predictor's guess from (h_t), and the encoder's answer once the observation is in. Part of its training loss is the gap between the two. Step 11 in the table above is that gap in miniature, where the model expected a shut door and the reading said otherwise. My model only ever has the answer, so when it forecasts blind its field stays wherever it was. It cannot imagine the door shutting at step 20, or opening in the first place. DreamerV3's guess is a distribution, so an imagined rollout can open the door in some samples and leave it shut in others.

The third is how the training is fed. The backward walk through each sequence is the same, but DreamerV3 cuts its experience into chunks of 64 steps drawn from a replay buffer that keeps filling while it acts, and its loss has more than one term: how well it reconstructs the observation, the reward and the continuation, plus two terms that pull the guess and the answer towards each other, weighted 0.5 and 0.1. Mine is one squared error over eighty fixed episodes.

The 101 post joked about representation collapse, where a model scores perfectly by declaring every state identical. That failure needs the model's target to be something the model produces itself. JEPA-style models, from LeCun's line of work, predict the representation of the next observation rather than the observation, so a representation that never changes predicts itself perfectly, and they keep a separate, slowly updated target encoder behind a stop-gradient to prevent it. My model's target is a thermometer reading, which it has no way to change. A constant field earns it nothing, and the model falls back to the four-weight fit and its 0.29. DreamerV3 has the same protection, because it reconstructs its observations, and a quieter risk from the other side: the term that pulls the encoder's answer towards the guess would, pushed far enough, teach the encoder to ignore the observation. DreamerV3 stops penalising that gap once it is below one nat, which it calls free bits.

The nearest thing this post has to collapse is a field that goes unused, and that is less final than it sounds. With every u weight at exactly zero the gradient is exactly zero, which is why they start small and random, but near zero it is small and still there. When the open-door leak came down to 0.10, the field stopped tracking the door because the signal was too faint to find in 4,800 rows, which is a shortage of log and not a trap the training fell into.

Try it

The playground has a scenario switch at the top. In the door scenario, recording adds the door, training runs the four-weight model and then the latent one, and the right-hand pane runs the room at minus four with a slider for where the five-step forecast starts. Drag that slider across step 10 and watch the orange forecast turn the wrong way. The open-door leak has a slider too; bring it down towards the closed rate and watch the memory weight in the latent model's table climb, then watch the field stop finding the door altogether.

The playground runs the same code with Math.random where the script's seed is, so your numbers will differ from the ones above by a little and agree with them in shape. The data behind each chart is linked under it.

The next post puts this model next to an LLM. One predicts the next token of text, the other predicts the next state of a room, and the interesting part is where an agent would want both.

Top comments (0)