A Path Towards Autonomous Machine Intelligence (LeCun, 2022) — part 7/12

Cleaned-text code points [84889, 98593). Bibliography omitted here; source.md includes it.
Pred1(s,a)  Pred1(s,a)
C(s[2])
 Pred1(s,a)  Pred1(s,a)
C(s[4])
Enc1(x)
s[1] s[2]
action
s[0]
Actor1
 Pred2(s,a)Enc2(s[0])
s2[0] s2[2]
Actor2
 Pred2(s,a)
C(s2[4])
s[3] s[4]
s2[4]
a2[2] a2[4]
a[1] a[2] a[3]a[0]
Figure 16: Hierarchical JEPA for Mode-2 hierarchical planning.
A complex task is deﬁned by a high-level cost computed from a high-level world-state representation
C(s2[4]). A sequence of high-level abstract actions (a2[2],a 2[4]) is inferred that minimizes C(s2[4]).
The inferred abstract actions are fed to lower-level cost modulesC(s[2]),C (s[4]) which deﬁne subgoals
for the lower layer. The lower layer then infers an action sequence that minimizes the subgoal costs.
Although only a 2-layer hierarchy is shown here, it is straightforward to extend the concept to multiple
levels.
The process described here is sequential top-down, but a better approach would be to perform a joint
optimization of the actions in all the layers.
(a2[2],a 2[4]) to optimize this objective. These high-level “actions” are not real actions but
targets for the lower level predicted states. One can think of them as conditions that the
lower-level state must satisfy in order for the high-level predictions to be accurate. Whether
these conditions are satisﬁed can be computed by cost modules C(s[2]) and C(s[4]). They
take a lower-level state s[2] and a high-level condition a2[2] and measure to what extent
the state satisﬁes the condition. With these subgoals deﬁned, the lower level can perform
inference and ﬁnd a low-level action sequence that minimizes the mid-level subgoals C(s[2])
and C(s[4]).
The process just described is top down and greedy. But one may advantageously iterate
the optimization so that high level and low-level action sequences are optimized jointly. The
cost modules may be conﬁgured by the conﬁgurator for the situation at hand.
The idea that an action is merely a condition to be satisﬁed by the level below is actually
an old one in control theory. For a example, a classical proportional servomechanism can be
seen as being given a target state. A quadratic cost measures the squared distance between
the target and the current state, and the control is simply proportional to the negative
gradient of the cost with respect to the action variables.
4.8 Handling uncertainty
The real world is not entirely predictable. Uncertainty in predictions of future world states
may be due to a number of reasons:
• the world is intrinsically stochastic (aleatoric uncertainty, type 1)
32

Pred1(s,a,z)  Pred1(s,a,z)
C(s[2])
 Pred1(s,a,z)  Pred1(s,a,z)
C(s[4])
Enc1(x)
s[1] s[2]
action
s[0]
a[1] a[2] a[3]a[0]
Actor1
 Pred2(s,a,z)Enc2(s[0])
s2[0] s2[2]
a2[2]
Actor2
 Pred2(s,a,z)
C(s2[4])
s[3] s[4]
s2[4]
a2[4]
z1[0]R1 z1[1]R1 z1[2]R1 z1[3]R1
z2[2]R2 z2[4]R2
Figure 17: Hierarchical JEPA for Mode-2 hierarchical planning in an uncertain environment.
Realistic environments are not entirely predictable, even when using highly-abstract representations.
Uncertainty about predictions can be handled by predictors with latent variables. The latent variables
(red circles) contain information about the prediction that cannot be derived from the prior observa-
tion. The latent variables must be regularized to prevent an energy collapse and to force the system
to predict as much as possible without the help of it.
At planning time, latent variables are sampled from distributions obtained by applying a Gibbs dis-
tribution to the regularizers. Each sample leads to a diﬀerent prediction. To produce consistent
latent sequences, the parameters of the regularizer can be functions of previous states and retrieved
memories.
As the prediction progresses, the number of generated state trajectories may grow exponentially. If
each latent variable has k possible discrete values, the number of possible trajectories will grow as
kt, where t is the number of time steps. Directed search and pruning strategies must be employed.
With multiple predicted trajectories, optimal action sequences can be computed that minimize the
average cost, or a combination of average and variance of the cost so as to minimize risk.
• the world is deterministic but chaotic, hence diﬃcult to predict without inﬁnitely
precise perception (aleatoric uncertainty, type 2)
• the world is deterministic but partially observable (aleatoric uncertainty type 3).
• the world is fully observable, but the sensors only give partial information about the
world state (epistemic uncertainty, type 1)
• the representation of the world state extracted by the perception module does not
contain the full information necessary for accurate prediction (epistemic uncertainty,
type 2).
• the world model is inaccurate due to limitations of its representational power (bounded
rationality or epistemic uncertainty, type 3).
33

• the world model is inaccurate due to it having been trained with limited amount of
data (epistemic uncertainty, type 4).
Much of the literature in reinforcement learning is focused on dealing with the stochastic
nature of the environment. It is often assumed from the start that models, critics and
policies must represent distributions. In the present work, we push the possible stochasticity
of a predicted variable into a latent variable, which may be optimized, predicted, or sampled.
This is what is often referred to in the ML literature as “the reparameterization trick”. We
do not need to use this trick here, since we view the latent-variable parameterization of the
predictions as fundamental.
Figure 17 represents a hierachical planning episode in the presence of uncertainty.
A prediction at a given level and time step, e.g. s2[2] requires a sample of the corre-
sponding latent variable z2[2]. The sample may come from the distribution whose negative
logarithm is the regularizer R2(z2[2])). The parameters of the regularizer may be constant
(e.g. ﬁxed Gaussian), predicted from currently-available data using amortized inference (e.g.
a multinomial or Gaussian whose parameters are computed from s2[0]) or produced by the
conﬁgurator. Using previous predictions to conﬁgure the latent regularizer biases the system
towards generating “good” trajectories.
As the prediction progresses, the number of generated state trajectories may grow ex-
ponentially: if each latent variable has k possible discrete values, the number of possible
trajectories will grow as kt, where t is the number of time steps. Directed search and prun-
ing strategies can be employed, as in classical Monte-Carlo Tree Search (MCTS). In the case
of continuous latents, one may sample latents from the continuous distributions deﬁned by
the regularizer.
Given a sample of all the latents, the optimal action sequences at every levels can be
inferred. However, the prediction process may need to be repeated for multiple drawings
of the latents, so as to cover the set of plausible outcomes. The inference process may be
used for multiple predictions to produce an action that does not just minimize the expected
cost, but also minimizes the uncertainty on the expected cost.
4.8.1 World Model Architecture
The details of the architecture of the world model should depend on the type of environment
the agent evolves in.
It is likely that the best module architectures in a JEPA should include some sort of
gating or dynamic routing mechanism.
For example, the best way to handle low-level, short-term predictions in videos is by
extracting simple local feature vectors and displacing those feature vectors from one frame
to the next, depending on predicted motions. The latent variables may encode a map of
displacements, which can modulate routing connections between one frame and the next.
For longer-term prediction at a higher level of abstraction, the relevant features are
objects and their interactions. The evolution may be best modeled by a transformer ar-
chitecture, which has the property of being equivariant to permutation and is appropriate
to capture interactions between discrete objects (Vaswani et al., 2017; Carion et al., 2020;
Battaglia et al., 2016).
34

Separating the World Model from the Ego Model: The natural world is complex
and somewhat unpredictable, requiring a powerful model with latent variables to account
for the unpredictability.
On the other hand, the agent itself is somewhat more predictable: a particular ac-
tion on eﬀector will produce a motion that can often be predicted deterministically. This
suggests that the agent should possess a separate model of itself, perhaps without latent
variables (Sobal et al., 2022) as the eﬀect of actions on proprioception somewhat easier to
predict than the evolution of the external world or the eﬀect of actions on it.
In turn, the ego-model that the agent has of itself can be used as a template for a model
of other agents in a multi-agent scenario.
4.9 Keeping track of the state of the world
Traditionally, modules in deep learning architectures communicate states through vectors
or multi-dimensional arrays. But this tends to be a very ineﬃcient method when the state
of the object being modeled only changes in minor ways from one time to the next.
A typical action of an agent will only modify a small portion of the state of the world.
If a bottle is being moved from the kitchen to the dining room, the states of the bottle, the
kitchen, and the dining room will be modiﬁed. But the rest of the world will be unaﬀected.
This suggests that the state of the world should be maintained in some sort of writable
memory. Whenever an event occurs, only the part of the world-state memory aﬀected by
the event is to be updated, while the rest is to be left unchanged.
A conventional key-value associative memory can be used for this purpose, similar to
what has been proposed in the context of memory-augmented networks (Bordes et al., 2015;
Sukhbaatar et al., 2015; Miller et al., 2016), and entity networks (Henaﬀ et al., 2017).
The output of the world model at a given time step is a set of query-value pairs (q[i],v [i]),
which are used to modify existing entries in the world-state memory, or to add new entries.
Given a query q, the world-state memory returns
Mem(q) =
∑
j
cjvj (15)
˜cj = Match( kj,q ) (16)
c = Normalize(˜c) (17)
where the kj are keys, the vj are stored values, function Match(k,q ) measures a divergence
or dissimilarity between a key and a query, vector c contains scalar coeﬃcients cj, and
function Normalize(˜c) performs some sort of competitive normalization or thresholding,
such as the commonly-used cj = exp(˜cj)/[γ+∑
k exp(˜ck)], where γis a positive constant.
Writing a valuer using query (or address) q into the memory can be done by updating
existing entries:
˜cj = Match( kj,q ) (18)
c = Normalize(˜c) (19)
vj = Update( r,vj,cj) (20)
Function Update(r,v,c ) may be as simple as cr + (1−c)v.
35

If the query is distant from all keys, the memory may allocate a new entry whose key is
q and corresponding value is r. The γconstant in the example Normalize function above
may serve as a threshold for acceptable key-query divergence.
One can view each entry as representing the state of an entity in the world. In the
above example of the bottle, the world model may contain keys kbottle,k kitchen,k dining−room
respectively representing the bottle, the kitchen and the dining room. The initial value of
vbottle encodes its location as “kitchen”, the inital value of vkitchen encodes its content as
including the bottle, and the initial value ofvdining−room encodes its content as not including
the bottle. After the event, the location and contents are updated.
All of these operations can be done in a diﬀerentiable manner, and would hence allow
to back-propagate gradients through them.
4.10 Data Streams
Much knowledge about the world is learnable through pure observation. The laws of mo-
tion of physical objects can, in principle, be derived from observation, without a need for
intervention. But training a world model eﬃciently may require more active or “agentive”
information gathering.
One can list ﬁve modes of information gathering with which an agent can learn about
how the world works:
1. passive observation: the agent is being fed a sensor stream (e.g. video, audio, etc)
2. active foveation: the agent is being fed a stream within which the focus of attention
can be directed without aﬀecting the environment. For example, watching a scene
while being able to orient the vision and sound sensors, or being being fed a wide-
angle, high resolution video and/or audio stream within which the focus of attention
can be directed.
3. passive agency: sensory streams in which another agent acting on the environment
is being observed, enabling the inference of causal eﬀects of agent actions on the state
of the environment.
4. active egomotion: the agent receives sensory streams from a real or virtual environ-
ment within which the position of the sensors can be modiﬁed without signiﬁcantly
aﬀecting the environment. This may include steerable active sensors (e.g. range sen-
sors, heat sensors, chemical sensors) as well as touch sensors.
5. active agency : sensory streams that are inﬂuenced by the agent’s actions. This
enables the establishment of causal models in which the agent can learn to predict the
consequences of its actions. This mode brings the exploration-exploitation dilemma
to the forefront.
In a complex environment, it may not be practical to collect enough passive data for the
world model to capture a suﬃcient portion of the environment’s behavior. Mode 2, 4, and
5 allow the agent to collect information that maximizes its understanding of the environ-
ment. But to do so may require intrinsic motivation modules that drive attention, curiosity,
36