Temporal Representation Alignment: Successor Features Enable Emergent Compositionality in Robot Instruction Following
Authors: Vivek Myers
Bill Zheng
Anca Dragan
Kuan Fang
Sergey Levine
Venue: NeurIPS 2025 poster ()
Keywords: Robotics
Representation Learning
Primary area: reinforcement_learning

Abstract: Effective task representations should facilitate compositionality, such
that after learning a variety of basic tasks, an agent can perform
compound tasks consisting of multiple steps simply by composing the
representations of the constituent steps together. While this is
conceptually simple and appealing, it is not clear how to automatically
learn representations that enable this sort of compositionality. We show
that learning to associate the representations of current and future
states with a temporal alignment loss can improve compositional
generalization, even in the absence of any explicit subtask planning or
reinforcement learning. We evaluate our approach across diverse robotic
manipulation tasks as well as in simulation, showing substantial
improvements for tasks specified with either language or goal images.

Temporal Representation Alignment: Successor Features Enable Emergent Compositionality in Robot Instruction Following
Vivek Myers* University of California, Berkeley
Bill Chunyuan Zheng*
Anca Dragan
Kuan Fangb
Sergey Levine
'Cornell University
Abstract
Effective task representations should facilitate compositionality, such that after learning a variety of basic tasks, an agent can perform compound tasks consisting of multiple steps simply by composing the representations of the constituent steps together. While this is conceptually simple and appealing, it is not clear how to automatically learn representations that enable this sort of compositionality. We show that learning to associate the representations of current and future states with a temporal alignment loss can improve compositional generalization, even in the absence of any explicit subtask planning or reinforcement learning. We evaluate our approach across diverse robotic manipulation tasks as well as simulation, showing substantial improvements for tasks specified with either language goal images.
1 Introduction
Compositionality is a core aspect of intelligent behavior, describing the ability to sequence previously learned capabilities and solve new tasks [1]. In domains involving long-horizon decision-making like robotics, various learning approaches have been proposed to enable this property, including hierarchical learning [2], explicit subtask planning [3, 4, 5], and
based "stitching" [6, 7]. In practice, these techniques are often unstable or data-inefficient in real-world robotics settings, making them difficult to scale [8].
By contrast, humans and animals are adept at quickly composing behaviors to reach new goals [1]. Possible explanations for these capabilities have been proposed, including the ability perform transitive inference [9], learn successor representations and causal models [10, 11]. and plan with world models [12]. In common among these theories is the idea of learning structured representations of the world, which inference about which actions will lead to certain goals.
How might these concepts translate algorithms for robot learning? In this work, study how adding an auxiliary successor representation learning objective affects compositional behavior in a real-world tabletop manipulation setting. We show that learning this representation structure improves the ability of the robot to perform long-horizon, compositionally-new tasks, specified either through goal images natural language instructions. Perhaps surprisingly, alignment does not need to he used for training the policy or test-time inference, as long as it is used auxiliary loss the same representations used for the tasks (Fig. 1).
found that this temporal
We compare our method, Temporal Representation Alignment (TRA), against past imitation and reinforcement learning baselines across set of challenging multi-step manipulation tasks in the Bridge setup [13] as well as the OGBench simulation benchmark [14]. Unlike prior work in setup, we focus the compositional capabilities of the robot policies: whole, the tasks
39th Conference on Neural Information Processing Systems (NeurIPS 2025).
Website: https: tttt://traappppp github.io/ "Equal contribution.

---

prior imitation and offline RI. methods
Instruction
Goal Image
"put food items bowl""
AWR
X
language E(t)
b(g)
representation
goal representation
temporal representation alignment
BC
policy x
TRA (OURS)
((s))
state representation
align state representations with tasks across time to compose behaviors
We show our Temporal Representation Alignment (TRA) method performing language task, "put all food items in the bowl." TRA adds a time-contrastive loss for learning task represen- tations to with goal- and Iagguage-condtiinee policy. While TRA can implicitly decompose the task into steps and execute them one by one, the behavioral cloning (BC) and offline RI. (AWR) methods fail this compositional task. The structured representations learned by TRA enable this compositional behavior without explicit planning or hierarchical structure.
Figure 1:
fit-of-dditributionn but each distinct subtask can be described through goal image that lies in the training distribution. Adding simple time-contrastive alignment loss improves compositional performance by >40% across 13 tasks 4 evaluation scenes, and simulation results show better per- formance compared to behavioral cloning (i.e., no structured representation learning), and comparable performance offline RL methods that explicitly use a learned value function.
2 Related Work
Our approach builds upon prior work on goal- and language-conditioned control, focusing particularly on the problem of compositional generalization.
Robot manipulation with language and goals. Recent improvements in robot learning datasets have enabled the development of robot policies that can be commanded with image goals and language instructions [5, 13, 15]. These policies can be trained with goal- and -enguagecooddiionned
imitation learning from human demonstrations [16, 17, 18. 19, 20], reinforcement learning [21, 22]. or other forms of supervision [23, 24]. When trained to reach goals, methods can additionally use hindsight relabeling [25, 26] improve performance [13, 27, 28, 29]. Our work shows how the benefits of soal-conditioned and language-conditioned supervised learning can be combined with temporal representation alignment to enable compositionality that would otherwise require planning or reinforcement learning.
Compositional generalization sequential decision making. the context of decision making. compositional generalization refers to the ability to generalize to new behaviors that are composed of known sub-behaviors [30, 31]. Biological learning systems show strong compositional generalization abilities [9, 28, 32, 33], and recent work has explored how similar capabilities can be achieved in artificial systems [34, 35, 36]. the context policy learning, exploiting the compositionality the behaviors can lead to generalization to unseen and temporarily extended tasks [6, 37, 38, 4, 39, 40].
Hierarchical and planning- partitioning task into its components [41, 42, 43, 44]. With improvements in vision-language models (VLMs), many recent works have explored using a pre-trained VLM to decompose task into subtasks that are more attainable for the low-level manipulation policy [5, 45, 46, 47, 42, 48, 49]. These approaches are limited by the need for robust pre-trained models that can be fine-tuned and prompted for embodied tasks. Our contribution is to show compositional properties can be achieved without any explicit hierarchical structure or planning, by learning a structured representation through time-contrastive representation alignment.
approaches also aim enable compositional behavior by explicitly

---

Representation learning for states and tasks. State and task representations for decision making aim improve generalization and exploit additional sources of data. Recent work in the robotics domain have explored the use of pre-trained representations across multimodal data, including images and language, for downstream tasks [50, 51, 52, 27, 53, 54, 55, 56, 57J. In reinforcement learning problems, representations are often trained to predict future states, rewards, goals, or actions [58, 59, 60, 61], and can improve generalization and sample efficiency when used as value functions [62, 63, 64, 65, 66]. Some recent works have explored the use of additional structural constraints on representations to enable planning [41, 43, 67, 68, 69], or enforced metric properties to improve compositional generalization [70, 71, 72].
The key distinction between our approach and past contrastive representation methods for robotics like VIP [59], GRIF [27]. and R3M [53] is that we focus on the
compositional generalization capabilities enabled by simply aligning representations across time addition to the task modalities, without using the learned representations for policy extraction or defining value function.
3 Temporal Representation Alignment
When training a series of short-horizon goal-reaching and instraction-folloiing tasks, our goal is to learn a representation space such that our policy can generalize new (long-horizon) task that can be viewed as a sequence of known subtasks. We propose to structure this representation space by aligning the representations of states, goals, and language in a way that is more amenable to compositional generalization.
Notation. We take the setting of a goal- and language-conditioned MDP M with state space S, continuous action space A (0, 1)da, initial state distribution Po. dynamics P(s' s, s,a), discount factor y, and language task distribution pe. A policy T(a s) maps states to a distribution over actions. We inductively define the k-step nuctio-connditonee) policy visitation distribution as:
pi 191, a1) - p(s1 s1,a)),
Pk+ ( a1) / / p(su+1 a) (s s1, s1, d s)
ph ( St. 01) _ p" (sk s),
(I)
Then, the discounted state visitation distribution can be defined as the distribution over st, the state reached after K ~ Geom(1 7) steps:
(s 1)
x (E |8,a).
(2)
R=0
We assume access to dataset of expert demonstrations D = K where each trajectory
T {St,s 1 ESXA
gathered by an expert policy B, and is then annotated with (r 81,1, 8H,)). Our aim is to learn policy that select actions conditioned we handle the continuous action space by representing both our policy and the expert policy as an isotropic Gaussian with fixed variance; we will equivalently write <(a 8, sp) denote the mode as à (1, sp) for a task sp.
language instruction E. As in prior work [13],
3.1 Representations for Reaching Distant Goals
We learn a goal-conditioned policy (a s, g) that selects actions reach goal g from expert demonstrations with behavioral cloning. Suppose we directly selected actions to imitate the expert on two trajectories in D:
81 + 82
8H
TED E
(4)
w + of + 8/1 - g
3

---

When conditioned with the composed goal we would be unable to imitate effectively as the composed state-goal (s.g) jointly out of the training distribution.
What would work for reaching g is to first condition the policy on the intermediate waypoint te, then upon reaching Wi, condition the goal g. the state-goal pairs (s,, (w, 9). and (5,, g) are all in the training distribution. If we condition the policy on some intermediate waypoint distribution p(w) (or sufficient statistics thereof) that captures all of these cases, we can stitch together the expert behaviors to reach the goal g.
Our approach is learn representation space that captures this ability,, that GCBC objective used in this space can effectively imitate the expert on the composed task. We begin with the goal-conditioned behavioral cloning [26] loss LBc
V. conditioned with waypoints w.
K
Lac (f as st 9: )
Σ
log π (a= 81, (g:)))
(5)
Enforcing the invariance needed to stitch Eq. (4) then reduces to aligning (b(g) + (w)) The temporal aligmment objective ((s) ++ o(s*) accomplishes this indirectly by aligning both ((w) and ()(g) to the shared waypoint representation ((w):
((()) K
1 LNCE ((ss; 1 1i
log
j= (o
+
Σ
log
i= colst (o(ss))))ss))
(6)
3.2 Interfacing with Language Instructions
To extend the representations from Section 1. to compositional instruction following with language tasks, we need some way to ground language into the va (future state) representation space. We use similar approach to GRIF [27], which uses an additional CL.IP-style [73] contrastive alignment loss with an additional pretrained language encoder E:
K
((E))
K
(a)) T(()))
CNCE
log
log
(7)
Σラ=
K
K 1==
3.3 Temporal Alignment
Putting together the objectives from Sections 3.1 and 3.2 yields the Temporal Representation Align- ment (TRA) approach. TRA structures the representation space of goals and language instructions to better enable compositional generalization. We learn encoders o, v. and & to map states, goals, and language instructions to a shared representation space.
f(u [ h(x:))
ENCE Di
Σ log
+ Σ log
(8)
k
K
((.))
Lec bac(fss, ot, ,tiKk 1) = Σ
log π (a: (s.E())) + log (a: sis ((.))())))
(9)
1=1
LTRA (1Ri st, Gi, k v.
(10)
N, -Lnc + 1 1i v. <)= + LxcR ((8,, 8 M di w) + LroF (kgr, li wi
behavioral cloning
temporal alignment
task alignment
Note that the NCE alignment loss uses CLIP-style symmetric contrastive objective [73, 67] highlight the indices in the NCE alignment loss (8) for clarity.
Our overall objective is to minimize Eq. (10) across states, actions, future states, goals, and language tasks within the training data:
min I. (e....e.....em..... an. [TRA (181,6s Ct,is Smin(ttk, 8H, () d, o. (11)
R.O.0.. (~Unnf((...

---

Algorithm I: Temporal Representation Alignment
1: input: dataset D ((8t,) 0t,
2: initialize networks e _ (T,, 0, ()
3: while training do
4: sample batch {(81,;; 0t,i; setk,s i (:)) D for k ~ Geom(1 77) 5: θ a VoLAA (st, 0t, Sethiss ( 1
6: output: language-conditioned policy x(ae (e: E(O)). goal-conditioned policy x = (g))
3.4 Implementation
summary of approach is shown in Algorithm In essence, TRA learns three encoders: ø. which encodes states, & which encodes future goals, and & which encodes language instructions. Contrastive losses used align state representations d(s;) with future goal representations (b(++,),, which are in turn aligned with equivalent language task specifications {(t) when available. We then learn behavior cloning policy ı that can be conditioned on either the goal or language instruction through the representation (r(g) or E()), respectively.
3.5 Temporal Alignment and Compositionality
We will formalize the intuition from Section 3.1 that TRA enables compositional generalization by considering the error on a "compositional" version of D, denoted D Using the notation from Eq. (3), we say D is distributed according to:
K
H
D D
II
(o((1.))p((f 81,3, 8H,) 1=1
II 1=1
Rh (ar, P(se41, 8t,i, at,),
(12)
or equivalently
DH
II (ol() 81,, 8H,) II (s,,)) P(3441, St.is 0r.),
(13)
by the isotropic Gaussian assumption. We will define D & Dl'' to be longer-horizon version of D extending the behaviors gathered under T across horizon oH > H' H that additionally satisfies a "time-isotropy" property: the marginal distribution of the states is uniform across the horizon, i.e., (o( (s(.).)) for all t {1...''}
will relate the in-distribution imitation error ERR( D) to the compositional out-of-distribution imitation error ERR( D*). We define
H
ERR (f; D) = ED
Σ 11
1 Ex [|āe, # sn.))l for ( 2, D. (14)
On the training dataset this is equivalent to the expected behavioral cloning loss from Eq. (9).
Assumption 1. The policy factorizes through inferred waypoints as:
goals: a(a s.9) x s, s,w) P(s: |
814k = dw
(15)
language: T(a) s, = / Taa s.w) P(st = w 8t-k = g) P(se+* = ()) dw dg.
(16)
where denote by (s,g) the MLE estimate of the action
Theorem 1. Suppose is distributed according to Eq. (12) and D is distributed according to Eq. (12). When 7 1 1/H and a > for optimal features 6 and under Eq. (11), have
ERR(T; n; D*) ERR(T; D)
a
1fa 2).
(17)
2a
2a
5

---

We can also define a notion of the language-conditinee compositional generalization error:
H
ERR ERR'(T; (T; -
ED
R [ =BE., 6)11
Corollary 1.1. Under the same conditions as Theorem 1,
( (u==² 1(a 2}. 2a
ERR (T: & ERR (T;;
The proofs well as a visualization of the bound are in Appendix F. Policy implementation details can be found in Appendix B
4 Experiments
We ask the following questions:
1. Can TRA enable zero-shot composition of tasks without additional rewards or planning? Does TRA improve compositional generalization over past methods?
3. How well does TRA capture skills that are less common within the dataset? temporal alignment by itself sufficient for effective compositional generalization?
Instruction Following Performance
Goal Reaching Performance
LCBC
GCBC
Octo
Octo
GRIF
GRIF
AWR
AWR
TRA(Ours)
TRA(Ours)
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Success Rate
Success Rate
(a) Language instruction tasks
(b) Goal-image conditioned tasks
Figure 2:
Aggregated performance on compositional generalization tasks, consisting of instruction- following and goal-reaching tasks.
4.1 Heal-World Experimental Setup
We evaluate TRA on collection of held-out compositionally-000 tasks tasks for which the in- dividual substeps are represented in the dataset, but the combination of those steps unseen. For example, in task such as "removing bell pepper from towel, and then sweep the towel", both the tasks "remove the bell pepper from the towel" and "sweep the towel" have similar entries within BridgeData, but such a combination of behaviors unseen. We utilize real-world robot manipulation interface with a DoF WidowX250 manipulator arm with 5Hz execution frequency. We train augmented version dataset [13], which contains over 50k
Swiruction
of the BridgeData trajectories with 72k language annotations. More details are in Appendix B.
Figure The tabletop manipulation setup used for the real-world evalua- tion (see 13).
In order to specifically test the ability of TRA to perform compositional generalization, we organize our evaluation tasks into 4 scenes that unseen in BridgeData, each with increasing difficulty:

---

Set A One-Step: These are the only tasks that are not compositionally-OO as all the tasks are one-step tasks. These tasks involve opening, putting an item in, and closing drawer, and have been seen in BridgeData,, although at lower frequency than object manipulation, and with new positions.
These tasks scene involves concatenating multiple tasks of the same nature in sequence, where robot must be able to perform all tasks within the same trajectory. During evaluation, we instruct the policy with instructions such as sweeping multiple objects in the scene that require composition (though are not sensitive to the order of the composition).
Set B Task Concatenation:
Set C Semantic Generalization: Unlike set (B, these tasks require manipulation with different objects of the same type. We test this using various food items within BridgeData, instructing the policy within a container.
Set D Tasks with Dependency: This is the most challenging set: these tasks have subtasks that require previous subtasks to be completed, such as taking an object out of a drawer.
The complete list of tasks is described in Appendix D.
4.2 Baselines
We compare against the following baselines in our real-world evaluation: GRIF [27] learns a goal- and language- conditioned policy using aligned goal image and language representations. In our experiments, this becomes equivalent to TRA when the temporal alignment objective is removed. GCBC [13] learns conditonn behavioral cloning policy that concatenates the goal image with the image observation. LCBC [13] learns a language-cone eeguuag--conitioned policy that concatenates the language with the image observation. OCTO [74] uses multimodal transformer to learn goal- and language-condinngg
policy. The policy is trained on the Open-X dataset [75], which incorporates BridgeData in its entirety. AWR [76] uses advantages produced by a value function to effectively extract a policy from an offline dataset. In our experiments, we use the difference between the contrastive loss between the current observation and the goal representation and the contrastive loss between the next observation and the goal representation as a surrogate for value function.
Table 1: Real-worlh Evaluation
Language-condtiooned
Goal-conditioned
Task
TRA
GRIF
LCBC
Octo
AWR
TRA
GRIF
GCBC
Octo
AWR
(A)
open the drawar 0.90 (-8.)
0.2y
(.60)t+13)) 0.60 (92)
0.40 (12))
0.60(+92)
0.40
0.80 (12)))
(A)
mushroom In drowar
0,50
0.g0 (tx2)
bagi (.4
0.90(+11)
0.49(++22)
0.50 (-30)
(A)
close drower 0.60
0.60
0.40
0.60
0.00
1.0(5500
(280
0 0.00
0.40(443)
(D)
take the from out of the drawer 0.60
0.0g
0.00
0,20
0.20
0.00(
0.40(:502)
0.00
(200
0.20(15112)
0,00
(B)
put the
towels (,80)9900
0.40 (Hroa)
0.29 (hi.))
1.00
0.40! +:02)
(100)
put the spoons on plates
tha0/
0 trea:
(C) put the com and sushi piato
0u0g
(140)
0.40
(.70
0.00/
0.20
(22)
0.30f
(C)
and mushroom howl
0.69
0.20/ 1921
(.50
0.20
0.4n! 1aa
(
(C) com. banama. and sushi in bowl
0.0g
(.20
0.50(
0.40
(D)
comn plato then sushi pptt
(170
0ugg(+oup (+o0g
0.40
0.60 ( (aa))
0.30(:+0))
(.20(++12)
(=000
(ha (100)
(A)
swoep to the right
(1889
0.40 (62)
0.40(++i2))
0.00
0.00
0.00f+10))
(B) fold clath inno the conner
1.00
0,29| (+z))
0.49
0.40
+92)
0,80(+01) (h4)
0:00
0:00-100
ousa
(B)
movs bell pepoer and twapp towel
(1.00
ougg
0.0g
0.20
0.20/
0.20
0.40
nooof
(A) One step tasks
Task concatenation
'hhe best-performing method(s) up to statistical significance are highlighted
(0) Semantic generulization
Tasks with dependency
We train GRIF, GCBC, LCBC, and AWR using the same augmented Bridge Dataset as TRA, and we use an Octo-Base 1.5 model for our evaluation. A more detail approach is detailed in Appendix C. During evaluation, we give all policies the same goal state and language instruction regardless of the architecture, as they are trained on the same language instruction with the exception of Octo, which doesn't benefit from paraphrased language data, but does benefit from more diverse language annotation set across larger datasct of varying length and complexity.
7

---

4.3 Real-world Evaluation
Our real-world evaluation aims to answer the following questions.
Does TRA enable compositionality? shows the success rates of the TRA method compared to other methods on real-world robot evaluation tasks. We marked all policies within the task orange if they achieve the best statistically significant performance, as determined by one-sided t-test with a significance level of 0.05 (see Appendix H for details). We first compare the performance against methods in A Although TRA performs well with drawer tasks, its performance against baseline methods is not statistically significant. However, TRA performs considerably better than that of any baseline methods on compositionally-001 instruction following tasks.
While TRA completed 88.9% of tasks seen in (B 83.3% of evaluations in Cn and 60% of tasks in D with instruction following, the best-performin baseline for (B was 30% with LCBC, 43.3% for with AWR, and 33.3% on with Octo. The same improvement was also present in goal reaching tasks, although at lower level, in which produced 60% success rate and scene D produced 43.3% success rate, as compared 46.7% and 20% for the best baselines.
How does TRA compare to conven- tional offline RL? While offline rein- forcement conventionally is considered necessary for "stitching" [77], demon- strate that TRA still outperforms offline reinforcement learning on robotic manip- ulation. TRA performs better than AWR for both language and image tasks, outper- forming AWR by 45% on instruction fol- lowing tasks, and by 25% on goal reaching tasks, showing considerable improvement offline RL method that promises compositional generalization via stitching.
Ablation: Using TRA as Value Signal
AWR+TRA
TRA (Ours)
AWR
AWR+TRA
TRA (Ours)
AWR
0.4
0.6
0.8
Success Rate
C Z Goal Images
Language
The policy trained with AWR often stops after one subtask, even though the goal in- struction image demanded all of the sub- tasks be completed. We see this in, e.g., Fig. 1, where 3 different policies use the same goal image for a task where all 3 food items must be put in the bowl.
Figure 4: Aggregated success rate of using AWR as an additional policy learning metric over all 4 scenes.
Does TRA help rarely-seen skills within the dataset? We also compare TRA against AWR across challenging tasks in the dataset. When conditioning on language, AWR struggles to effectively generalize to compositionally harder tasks, with average success rate decreasing from 43.3% in to 6.67% from C to D. compared to a decrease of only 83.3% to 60% for TRA. Other agents do not perform well as AWR in ( as the lack of such compositional generalization prevented the policies from achieving all of the tasks at a reliable rate.
Is TRA sufficient in achieving compositional generalization? We demonstrate in our real- world experiment that only using temporal alignment is sufficient for achieving good compositional generalization. We evaluate this by comparing a policy trained on only temporal alignment loss (our method), and another policy trained such loss and have these losses weighed by AWR. Our AWR implementation detailed in Appendix C.
Figure 4 shows that across all evaluation tasks, AWR provides no additional benefit on top of temporal alignment. In fact, using AWR marginally decreases the efficacy of TRA, unlike when used with GCBC and LCBC.
8

---

Table 2: OGBench Evaluation
Methods
Task
TRA
GCBC
CRL
GCIQL
GCIVL
QRL
60,7(*30
45,5(#39)
52,76=22)
29.3
(4.1
(924)
antraze large atitch
12.8
3,4(110)
(0.8(49))
7.5()07)
18,5
(8.4))97)
stitch
17,0(211)
24.5(222)
0,7
2.1 (,10)
21.42299
(.8(2
46.1 (6.19)
29.0(*17)
36,2(#99)
(21a))
(2,3(+96)
(8.0(
8.6(214)
5.6(210)
4.0(299))
0,5
(22()9))
3,5(293)
bunanol draze large stitch
35.4(211)
24.0(2000
82.8(414)
34.2(211)
15,7(219)
74.6(2)))
large navigate
cube single soiay
(,2()
8,4
38.3
99.3(283)
70.6(430)
(5,5(221)
RL methods with separate value network to update the actor in gray.
The best non-RL. methods up to significance are highlighted.
1 We bold best performance.
4.4 Testing Compositionality in Simulation
We also tested the compositional behavior in simulation using tasks from the OGBench [14] offline RL benchmark suite. This environment features environments for locomotion and manipulation, each with multiple offline datasets that can be used for training, including one that explicitly tests compositional generalization (the "stitch" datasets) by creating multiple short datasets that comprise single, larger task. Some environments can be seen in Fig. 5 We modify TRA to account for the lack of language instructions. See Appendix for details.
We evaluate the performance of TRA on seven different environments in OGBench. In of these environments we use the "stitch" dataset, while two other environ- ment use a more general goal-reaching dataset ("navigate" and "noisy"). Table 2 shows the performance of TRA compared to other non-hierarchical methods on these environments from OGBench. Consis- tent with our real-world results Table and Fig. 4, TRA outperforms other imi- tation and offline RL methods on certain environments that require compositional generalizations, including CRL [78] that also has separate value and critic network.
Figure 5: Two environments from the OGBench suite [14]. Left: a cube stacking environment. Right: a humanoid maze navigation environment.
In non-stitching environments, while traditional offline RL methods outperform TRA, TRA still improves GCBC.
4.5
Failure Cases
Since we use a Gaussian policy, TRA struggles when multimodal behavior is observed, and sometimes fails to reach the goal due to early grasping or incorrect reaching [37]. While TRA did seem to provide small improvements the in-distribution tasks of A, the primary benefits derived from TRA were seen on compositionally-OOD tasks. We further discuss failure cases in Appendix E.
5 Conclusions and Limitations
In this paper, we studied a temporal alignment objective for the representations used in (goal- and behavior cloning. This additional structure provides robust compositional generalization capabilities in both real-world robotics tasks and simulated RL benchmarks. Perhaps surprisingly, these results suggest that generalization properties usually attributed to reinforcement learning methods may be attainable with supervised learning with stal-structured, temporally- consistent representations.

---

Limitations and Future Work While TRA consistently outperformed behavior cloning in real world and simulation evaluations, the degree of improvement degrades when behavior cloning cannot solve the task at all. Future work could examine how to improve compositional generalization in such cases through additional structural constraints on the representation space. To scale to more complex settings, similar approaches with more complex architectures such as transformers and diffusion policies may be needed for policy and/or representation learning. Methods like TRA that learn compositional task representations could be used with more complex models like VLMs [73], VLAs [79], or LLMs [80] to improve their generalization capabilities. In these settings, future work should examine safety of task generalization when interacting with humans and avoid emergent misalignment [81].
References
[1] K. S. Lashley The Problem of Serial Order in Behavior. Cerebral Mechanisms in Behavior, pp. 112-136. 1951.
[2] Tejas D Kulkarni, Karthik Narasimhan, Ardavan Saeedi, and Josh Tenenbaum. Hierarchical Deep Reinforcement Learning: Integrating Temporal Abstraction and Intrinsic Motivation. Neural Information Processing Systems, volume 29, 2016.
[3] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and Offline Reinforcement Learning by Planning With Learned Model. Neural Information Processing Systems, volume 34, pp. 27580-27591 2021.
[4] Kuan Fang, Patrick Yin, Ashvin Nair, Homer Walke, Gengchen Yan, and Sergey Levine. Gener- alization With Lossy Affordances: Leveraging Broad Offline Data for Learning Visuomotor Tasks. Conference on Robot Learning, 2022.
[5] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, et al. Do I Can, Not as I Say: Grounding Language in Robotic Affordances. Conference on Robot Learning, 2022.
[6] Raj Ghugare, Matthieu Geist, Glen Berseth, and Benjamin Eysenbach. Closing the Gap Between TD Learning and Supervised Learning Generalisation Point of View. International Conference on Leurning Representations, 2023.
[7] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline Reinforcement Learning With Implicit Q-Learning. International Conference on Learning Representations, 2022.
[8] Cassidy Laidlaw, Banghua Zhu, Stuart Russell, and Anca Dragan. The Effective Horizon Explains Deep RL Performance in Stochastic Environments. International Conference on Learning Representattonss, 2024.
[9] Simon Ciranka, Juan Domdigo, Ivan Padezhki, Clara Wicharz, Charley M. Wu, and Bern- hard Spitzer. Asymmetric Reinforcement Learning Facilitates Human Inference of Transitive Relations. Nature Human Behaviour, 6(4):555-564 2022.
[10] Peter Dayan. Improving Generalization for Temporal Difference Learning: The Successor Representation. Neural Computation, volume 5, pp. 613-624, 1993.
[11] Alison Gopnik, Shaun O'Grady, Christopher G. Lucas, Thomas I. Griffiths, Adrienne Wente, Sophie Bridgers, Rosie Aboody, Hoki Fung, and Ronald E. Dahl. Changes in Cognitive Flexibility and Hypothesis Search Across Human Life History From Childhood to Adolescence to Adulthood. National Academy of Sciences, 11430))7892-7899, 2017.
[12] Oliver M. Vikbladh, Michael R. Meager, John King, Karen Blackmon, Orrin Devinsky, Daphna Shohamy, Neil Burgess, and Nathaniel D. Daw. Hippocampal Contributions to Model-Based Planning and Spatial Memory. Neuron, 102(3))683-692 2019.
[13] Homer Rich Walke, Kevin Black, Tony 7.. Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen- Estruch, Andre Wang He, Vivek Myers, et al. BridgeData V2: A Dataset for Conference on Robot Learning at Scale. Conference on Robot Learning, pp. 1723-1736, 2023.
[14] Seohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. OGBench: Bench- marking Offline Conaiioond RL. International Conference on Learning Representations, 2025.
[15] Mohit Shridhar, Lucas Manuelli, and Dieter Fox. CLIPort: What and Where Pathways for Robotic Manipulation. Conference on Robot Learning, 2021.
10

---

[16] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, et al. PaLM: Scaling Language Modeling With Pathways. J. Mach. Learn. Res, 2023.
[17]
Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan. VIMA: General Robot Manipulation With Multimodal Prompts. International Conference on Machine Learning, 2023.
[18]
Corey Lynch and Pierre Sermanet. Language Conditioned Imitation Learning Over Unstructured
Data, Robotics: Science and Systems XVII, 2021.
[19] Corey Lynch, Ayzaan Wahid, Jonathan Tompson, Tianli Ding, James Betker, Robert Baruch, Travis Armstrong, and Pete Florence. Interactive Language: Talking to Robots in Real Time.
IEEE Roboties and Automation Letters, (arXiv:221000440)):
2023.
[20]
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choro- manski, Tianli Ding, Danny Driess, al. RT-2: Vision-Language-Actionn Models Transfer Web Knowledge to Robotic Control. Conference on Robot Learning, 2023.
[21] Yevgen Chebotar, Quan Vuong, Karol Hausman, Fei Xia, Yao Lu, Alex Irpan, Aviral Kumar, Tianhe Yu, et al. Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions. Conference on Robot Learning, 2023.
[22] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. 2021.
[23] Andreea Bobu, Yi Liu, Rohin Shah, Daniel S. Brown, and Anca D. Dragan. SIRL: Similarity- Based Implicit Representation Learning. ACM/IEEE International Conference on Silman-RRobn Interaction, pp. 565-574, 2023.
[24] Yuchen Cui, Siddharth Karamcheti, Raj Palleti, Nidhya Shivakumar, Percy Liang, and Dorsa Sadigh. No, to the Right: Online Language Corrections for Robotic Manipulation via Shared Autonomy. ACM/IEEE International Conference on Stmann-Robot Interaction, pp. 93-101, 2023.
(25) Marein Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight Experience Replay. Neural Information Processing Systems, volume 30, 2017.
[26] Leslie Pack Kaelbling. Learning to Achieve Goals. International Joint Conference on Artificial Intelligence, 1993.
[27] Vivek Myers, Andre Wang He, Kuan Fang, Homer Rich Walke, Philippe Ching-An Cheng, Mihai Jalobeanu, Andrey Kolobov, Anca Dragan, and Sergey Levine. Goal Representations for Instruction Following: A Semi- Semi-Supervised Language Interface to Control. Conference on Robot Learning, pp. 3894-3908, 2023.
[28] Stanislas Dehaene, Fosca Al Roumi, Yair I Lakretz, Samuel Planton, and Mathias Sablé-Meyer. Symbols and Mental Programs: Hypothesis About Human Singularity. Trends in Cognitive Sciences, 26(9):751-766, 2022.
[29] Yiming Ding, Carlos Florensa, Picter Abbeel, and Mariano Phiclipp. Coul-Conditioned Imitation
Learning. Neural Information Processing Systems, 32, 2019.
[30] Valerio Rubino, Mani Hamidi, Peter Dayan, and Charley M. Wu. Compositionality Under Time Pressure. Cognitive Science Society, volume 45, 2023.
[31] Mark Steedman. Where Does Compositionality Come From? AAAI Technical Report, 2004.
[32] David W. Dickins. Transitive Inference in Stimulus Equivalence and Serial Learning. European Journal of Behavior Analysis, 12(2):523-555, 2011.
[33] Brenden M. Lake, Tal Linzen, and Marco Baroni. Human Few-Shot Learning of Compositional Instructions. CogSci, 2019.
[34] Ekin Akyürek, Afra Feyza Akyürek, and Jacob Andreas. Learning to Recombine and Resample Data for Compositional Generalization. International Conference on Learning Representations, 2021.
[35] Takuya Ito, Tim Klinger, Doug Schultz, John Murray, Michael Cole, and Mattia Rigotti. Compositional Generalization Through Abstract Representations in Human and Artificial Neural Networks. Neural Information Processing Systems, 35:32225-322333 2022.
11

---

[36] Martha Lewis, Nihal V. Nayak, Peilin Yu, Qinan Yu, Jack Merullo, Stephen H. Bach, and Ellie Pavlick. Does CLIP Bind Concepts? Probing Compositionality in Large Image Models. Conference of the European Chapter of the Association for Computational Linguistics, 2024.
[37] Aviral Kumar, Anikait Singh, Frederik Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre- Training for Robots: Offline RL Enables Learning New Tasks in Handful of Trials. Robotics: Science and Systems XIX, 2023.
[38] Kuan Fang, Yuke Zhu, Animesh Garg, Silvio Savarese, and Li Fei-Fei. Dynamies Learning With Cascaded Variational Inference for Multi-Step Manipulation. Conference on Robot Learning, 2019.
[39] Ajay Mandlekar, Danfei Xu, Roberto Martín-Martín, Silvio Savarese, and Li Fei-Fei. Learning to Generalize Across Long-Horizon Horizon Tasks From Human Demonstrations. arXiv:2003.0600 06085,
2021.
[40] Soroush Nasiriany, Vitchyr H. Pong, Steven Lin, and Sergey Levine. Planning With Goal- Conditioned Policies. 10Xiv:1911.08553 2019.
[41] Kuan Fang, Patrick Yin, Ashvin Nair, and Sergey Levine. Planning to Practice: Efficient Online Fine-Tuning by Composing Goals in Latent Space. International Conference on Intelligent Robots and Systems, 2022.
[42] Vivek Myers, Bill Chunyuan Zheng, Oier Mees, Sergey Levine, and Kuan Fang. Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation. Conference on Robot Learning, 2024.
[43] Tianjun Zhang, Benjamin Eysenbach, Ruslan Salakhutdinov, Sergey Levine, and Joseph E Gon- zalez. C-Planning: An Automatic Curriculum for Learning Goal- Real-Re Tasks. International Conference on Learning Representations, 2022.
[44] Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. HIQL: Offline Goal- Conditioned RL With Latent States as Actions. Neural Information Processing Systems, 2023.
[45] Maria Attarian, Advaya Gupta, Ziyi Zhou, Wei Yu, Igor Gilitschenski, and Animesh Garg. See, Plan, Predict: Language-Guided Cognitive Planning With Video Prediction. arXiv:2210.0382, 2022.
[46] Suneel Belkhale, Tianli Ding, Ted Xiao, Pierre Sermanet, Quon Vuong, Jonathan Tompson, Yevgen Chebotar, Debidatta Dwibedi, and Dorsa Sadigh. RT-H: Action Hierarchies Using Language. arXiv:2403.00 0.144000.112,,... 2024.
[47] Minae Kwon, Hengyuan Hu, Vivek Myers, Siddharth Karamcheti, Anca Dragan, and Dorsa Sadigh. Toward Grounded Commonsense Reasoning. International Conference Roboties and Automation, 2023.
[48] Ishika Singh, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. ProgPrompt: Generating Situated Robot Task Plans Using Large Language Models. International Conference on Robotics and Automation, 2023.
[49] Zichen Zhang, Yunshuang Li, Osbert Bastani, Abhishek Gupta, Dinesh Jayaraman, Yecheng Ja- son Ma, and Luca Weihs. Universal Visual Decomposer: Long-Horizon Manipulation Made Easy. arXiv:2310.085 2023.
[50] Siddharth Karamcheti, Suraj Nair, Annie Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language- Driven Representation Learning for Robotics. Robotics Science and Systems, 2023.
[51] Liunian Harold Li, Pengchuan Zhang, Haotian Zhang, Jianwei Yang, Chunyuan Li, Yiwu Zhong, Lijuan Wang, Lu Yuan. Lei Zhang, Jenq-Neng Hwang, Kai-Wei Chang, and Jianfeng Gao. Grounded Language- Image Pre- Traininng IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10955-10965. 2022.
[52] Yecheng Jason Ma, William L Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-Iman Image Representations and Rewards for Robotic Control. International Conference on Machine Learning, 2023.
[53] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A Universal Visual Representation for Robot Manipulation. Conference on Robot Learning, pp- 892-909, 2022.
12

---

[54] Jyothish Pari, Nur Muhammad (Mahi) Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The Surprising Effectiveness of Representation Learning for Visual Imitation. Roboties: Science and Systems XVIII, 2022.
[55] Rutav Shah and Vikash Kumar. RRL: Resnet as Representation for Reinforcement Learning. International Conference on Machine Learning, 2021.
[56] Yuchen Cui, Scott Niekum, Abhinay Gupta, Vikash Kumar, and Aravind Rajeswaran. Can Foundation Models Perform Zero-Shot Task Specification for Robot Manipulation? LADC, 2022.
[57] Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler. Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. BC-Z: Zero-Shot Task Generalization With Robotic Imitation Learning. Conference on Robot Learning, p. 12, 2021.
[58] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R. Devon Hjelm. Unsupervised State Representation Learning Atari. Neural Information Processing Systems, 2019.
[59]
Yecheng Jason Ma. Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training. International Conference on Learning Representations, 2023.
[60] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning Invariant Representations for Reinforcement Learning Without Reconstruction. International Conference on Learning Representations, 2021.
[61] Linxi Fan, Guanzhi Wang. Yunfan Jiang, Ajay Mandlekar, Yuncong Yang. Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar. MineDojo: Building Open-Ended Embodied Agents With Internet-Scale Knowledge. Neural Information Processing Systems, 2022.
[62] André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Has- selt, and David Silver. Successor Features for Transfer in Reinforcement Learning. Neural Information Processing Systems, volume 30, 2017.
[63] Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning Successor States and Goal- Dependent Values: A Mathematical Viewpoint. arXiv:2101.0711 123,,, 2021.
[64] Peter Dayan. Improving Generalisation for Temporal Difference Learning: The Successor Representation. Neural Computation, 1993.
[65] Alexey Dosovitskiy and Vladlen Koltun. Learning Act by Predicting the Future. International Conference on Learning Representations, 2017.
[66] Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Varia- tional Empowerment as Representation Learning for Goal-Conditioned Reinforcement Learning. International Conference on Machine Learning, pp. 1953-1963, 2021.
[67] Benjamin Eysenbach, Vivek Myers, Ruslan Salakhutdinov, and Sergey Levine. Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference.
arXiv:2403.0444 004022, 2024.
[68] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamies for Planning From Pixels. 1eXiv:1811.044 2019.
[69] Vivek Myers, Catherine Ji, and Benjamin Eysenbach. Horizon Generalization in Reinforcement Learning- International Conference on Learning Representations, 2025.
[70] Bo Liu, Yihao Feng, Qiang Liu, and Peter Stone. Metric Residual Network for Sample Efficient Goal-Conditioned Reinforcement Learning. AAAI Conference on Artificial Intelligence, volume 37. pp. 8799988860 2023.
[71] Vivek Myers, Chongyi Zheng, Anca Dragan, Sergey Levine, and Benjamin Eysenbach. Learn- ing Temporal Distances: Contrastive Successor Features Can Provide Metric Structure for Decision-Making. International Conference on Machine Learning, #xXiv22406.17098 2024.
[72] Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal Reinforcement Learning via Quasimetric Learning. International Conference on Machine Learning, pp. 36411 1-36430, 2023.
[73] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya
13

---

Sutskever. Learning Transferable Visual Models From Natural Language Supervision. Interna- tional Conference on Machine Learning, arXiv:2103.0000 2021.
[74]
Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, al. Octo: An Open-Source Generalist Robot Policy. Robotics: Science and Systems, 2024.
[75] Abby O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lec, Acorn Pooley, Agrim Gupta, et al. Open X- X-Embodiment: Robotic Learning
Datasets and RT-X Models. International Conference on Roboties and Automation, 2024.
[76] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage- Weighted Regres- sion: Simple and Scalable Off-Policy Reinforcement Learning. arXiv: 1910.00177, 2019.
[77] Aviral Kumar, Joey Hong. Anikait Singh, and Sergey Levine. Should I Run Offline Reinforce- ment Learning or Behavioral Cloning? International Conference on Learning Reprexentations,
2021.
[78] Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive Learning Coal-Conditioned Reinforcement Learning. Neural Information Processing Systems, 35:35603-3566 -30000, 2022.
[79] Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, et al. OpenVLA: An Open-Source Vision-Language-Are Model. arXiv:2406.09244 092466 2024.
[80] Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, et al. Language Models Are Few-Shot Learners. Neural
Information Processing Systems, 2022.
[81] Simon Zhuang and Dylan Hadfield-Menell. Consequences of Misaligned AI, Neural Information Processing Systems, volume 33, pp. 15763-15773, 2020.
[82] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. International Conference Learning Representations, 2015.
[83] Gerhard Neumann and Jan Peters. Fitted Q-Iteration by Advantage Weighted Regression.
Neural Information Processing Systems, volume 21, 2008.
A Code and Website
An implementation of TRA is available at https://anonymous. www.://muuu.......oooo... science/ei/ ogcr1-4344/.
B TRA Implementation
In this section, provide details on the implementation of temporal representation alignment (TRA) and its training process.
B.1 Dataset Curation
We use an augmented version of BridgeData. We augment the dataset by rephrasing the language annotations, as described by [27], with 5 additional rephrased language instruction for each language instruction present in the dataset, and randomly sample them during training.
During data loading process, for each observation that is sampled with timestep k. also sample k+ A min(k + xr, H), Geom(1 - 7), and load 8& along with 8*+. We employ random cropping, resizing, and hue changes during training process image robustness. We set 7 = 0.95 for policy training on BridgeData.
B.2 Policy Training
We use a ResNet-34 architecture for the policy network. We train our policy with one Google V4-8 TPU VM instance for 150,000 steps, which takes total of 20 hours. We use learning rate of 3 x 10-4 with the ADAM optimizer [82], 2000 linear warm-up steps, and a MLP head of 3 layers of 256 dimensions after encoding the observation representations as well as goal representations.
14

---

C
Baseline Implementations
We summarize the implementation details of the baselines discussed in Section 4.2.
C.1
Octo
We use the Octo-base 1.5 model publicly available on HuggingFace for evaluating Octo baselines. We inference code that is readily available for both image- and language- conditioned tasks. During inference, we use an action chunking window and an execution horizon window 4.
C.2 Behavior Cloning
We use the same architecture for LCBC and GCBC as in Walke et al. [13], Myers et al. [27]. During the training process we use the same hyperparameters as TRA.
C.3 Advantage Weighted Regression
In order to train an AWR agent without separately implementing reward critic, we follow Eysenbach al. [78] and use a surrogate for advantage:
A(si) ExCe (f(ss). f(g)) (xc=(((((+))
(18)
Here, f can be any of the encoders Φ, &, th. is the same InfoNCE loss defined Section 3, and g defined as either the goal observation the goal language instruction, depending the modality.
And extract the policy using advantage weighted regression (AWR) [83]:
arg max E, B, log n (als s, exp(A(s.a)/3))
(19)
During training. we set ( to 1, and we use batch size of 128, the same value as policy training for our method.
D Experiment Details
In this section, we go through our experiment details and how they are set up. During evaluation, randomly reset the positions of each item within the table, and perform to 10 trials on each task, depending whether this task is important within each scene. We examine tasks that seen in BridgeData, which include conventionally less challenging tasks such as object manipulation, and challenging tasks to learn within the dataset such cloth folding and drawer opening.
D.1
List of Tasks
Table describes each task within each scene, and the language annotation used when the policy is used for inference. Every task that is outside of the drawer scene are multiple step, and require compositional generalization.
D.2 Inference Details
During inference, we use a maximum of 200 timesteps to account for long-horizon behaviors, which remains the same for all policies. We determine task as successful when the robot completes the task it was instructed to within the timeframe. For evaluating baselines, we use trials for each of the tasks.
E Additional Visualizations
In this section, we show additional visualizations of TRA's execution compositionally-OO tasks. We use folding, taking mushroom out of the druwer, and corn on plate, then sushi in the pot as examples which require compositionality Fig, 6 and Appendix E.
15

---

Table 3: Task Instructions
Scene
Count
Task Description
Instruction
10
open drawer
drowel"
Drawer
10
the mushroom the drawer
the drawer"
10
close the drawer
"closa the drower"
10
the spoons on the plates
"mooe the sprion onto the pirtles.
Generalization
10
put the spoons on the towels
"move the spocns on the towes"
10
fold the cloth into the center from all corners sweep the towels to the right
"fokd the ctoth info cenneer"
10
"twisep the
right
10
put the sushi and the corn on the plate
"put the food htoms on the plate"
Semantic Generalization
10
put the sushi and the mushroom the howl
'pul the food items in the bow!
10
put the sushi, corn, and the hamana in the bowl
'put evarything the
mushroom
"open the drowor and than take the mush-
"move the ball pepper to the bottom right
Tasks With Dependency
10
move hell pepper and sweep towel
corner the toble and men sweeD
10
the carn the plate, and then pat the sushi in the
the susti in the pof"
"move the bell pepper to the bottom right of the table, and then move the towel the top right the table"
LCBC
TRA (OURS)
Figure Example rollouts of a task with TRA and LCBC. While TRA is able to successfully compose the steps to complete the task, LCBC fails to ground the instruction correctly.
E.1
Failure Cases
We break down failure cases in this section. While TRA performs well in compositional generalization, cannot counteract against previous failures seen with behavior cloning with a Gaussian Policy.
Analysis of Compositionality
We prove the results from Section 3.5.
F.1 Goal Conditioned Analysis
Theorem 1. Suppose D is distributed according to Eq. (12) and D* is distributed according to Eq. (12). When 7 1/H and a >> for optimal features & and under Eq. (11), have
a- +
( (a (a--2 2
ERR(T; D*) < ERR(R;; D)
1fa > 2}.
(17)
2a
2a
Proof. We have from Eq. (14) for K ~ Geom(1 - 7):
H' lā. (8r...8))l
BRR(TT;"") A ED.
A
16

---

"fold the towel into center"
"open the drawer, and then take the mushroom out of the drawer"
"put the corn on the plate, and then put the sushi in the pot"
Figure TRA performs compositional generatlization over a variety of tasks seen within BridgeData.
"open the drawer, and then take the mushroom out the drawer"
"put everything the bowl"
Figure 8: Most of the failure cases came from the fact that a policy cannot learn depth reasoning, causing early grasping late release, and it has trouble reconciling with multimodal behavior.
17

---

H' -2H
H'-H
ED
δd, x (i
+
ED-
Σ
H'
Ad
H'
ndA
H' 2'-22H+1
H¹
E
11'
Ad
1= HH+1
-
S
D
|8,,1 (Bc.i.8il)
RD
Σ
H
AdA
H
nd
t=1l' H++1
1=
21
a - 2
1/a > 27
20
15 ,311
S
D
H'
Add
t=1l' H++1
11'-11 11
l1 p K..:1& [] + [0]]+==-_)) (a-3)
ED
Σ Ex
1fa > 2]
H
ndd
2a
I=H 22
H'
x Ber 2012
E D
H'
Ad
t=
HH+1
H'
1
E
Σ
Ex
lla
(81.11#'' KK
a (a--2)
1 2]
H
Add
2a
2H
lde,
(Be. 111
I H
D
Add
t=Hl'
H++1
H'-H
ED
Σ
Ex
|1
(R.. (
+
(a-2) 2
1fa > 2]
H
ndd
2a
1= -2H
- g
+ (a-2)
< ERR(T; (eR(T:D")=> + ED
1fa > 2}
H
1 1 7
2a
(RR(T:D")) a -
2
+
1fa 2}.
(20)
20
2(x
L
F.2
Language Conditioned Analysis
Corollary 1.1. Under the same conditions as Theorem 1,
ERR' (T: D*) < ERR' (T; D)
a
+ (a- 1fa > 2}.
2a
2a
The proof is similar to Appendix F.1, but over the predictions of instead of sb.
F.3 Visualizing the Bound
We compare the bound from Theorem with the "worst-case" bound of (RR(TT; D*) ERR(T; D) in Fig. 9. The bound from Theorem 1 is tighter than the worst-case bound, and shows that the compositional generalization error decreases as increases.
G OGBench Implementation Details
To implement TRA in OGBench, which does not have a corresponding language label all goal- reaching tasks, make the following revision to TRA to accommodate the lack of a language task.
18

---

Compositional Generalization Error Bound
ERR (RRR(TDD)) (T;")) ERR (T, D)
0.6
0.4
0.2
1
1.2
1.4
1.6
1.8
2.2
2.4
Q
bound (17)
worsE case
Figure 9: Visualizing the bound (Eq. 17 from Theorem 1) the compositional generalization emror.
Table 4: Success Rate for Different GCBC Architectures in OGBench.
Environment
GCBC
GCBC-6
antmaze medium stitch
45.5
48.7 48.7 (2.7)
antmaze large stitch antsoccer arena stitch
3.4+(1.0)
6.8+(13)
24.5
1.4+
humanoidmaze medium stitch
29.04(1.7)
34.4-(1.7)
humanoidmaze large stitch
5.6=(1.0)
3.5+(1.1)
antmaze large navigate cube single noisy
24+(0.6)
16.1
8.4=(1.0)
8.7+(0.))
Table 5: TRA hyperparameters.
Hyperparamecte
Value
State and goal encoder dimensions State and goal encoder latent dimension
(64, 64, 64)
64
Discount factor
0.995 (large locomotion environments), (other)
Alignment coefficient a
60 (medium Incomotion), 100 (large locomotion), 20 (non-stitch)
19

---

Table 6: Statistical comparisons of language-conditoon success rates
t-values
p-values
Task
Best Score
TRA GRIF
LCBC Octo
AWR
TRA
GRIF
LCBC
Octo
AWR
(A)
open
0.80
0.00
2.50
0.72
0.72
1.43
0.500
0.038
0.498
0.498
0.198
mushroom drower
0.80
0.00
0.00
1.43
6.00
0.72
0.500
1.000
0.198
<0.001
0.498
closs drower
0.60
0.00
0.00
0.58
0.00
0.58
0.500
1.000
0.580
1.000
0.580
item out of drawer
0.60
0.00
3.67
3.67
1.55
3.67
0.500
0.005
0.005
0.155
0.005
(B) put the spcons on towct
0.80
0.00
1.43
2,50
6.00
2.50
0.500
0.198
0.038
<0.001
0.038
the
the platos
0.90
0.00
3.13
3.13
3.13
9.00
0.500
0.020
0.020
0.020
<0.001
(C) put the corn and sushi on plate
0.90
0.00
9.00
1.89
9.00
2.06
0.500
<0.001
0.113
<0.001
0.058
suhi and mushroom
0.80
0.00
6.00
0.72
2.50
0.95
0.500
<0.001
0.498
0.038
0.356
(C) com. banana, sushi
0.80
0.00
6.00
6.00
6.00
3.18
0.500
<0.001
<0.001
<0.001
0.005
on plate
pot
0.70
0.00
4.58
1.04
0.35
1.99
0.500
0.001
0.332
0,739
0.079
(A)
sweep to the fight
0.80
0.00
2.50
1.43
1,43
6.00
0.500
0.038
0.198
0.198
<0.001
the center
1.00
0.00
4,00
2.45
2.45
2.45
0,500
0.016
0.070
0.070
0.070
move pepper sweep towel
0.50
0.00
3.00
3.00
1.15
3.00
0.500
0.015
0.015
0.277
0.015
For cach task the reported I-values compare the candidate method to the best-performing method. We bold the t/p-values corresponding methods that do not significantly (p 0.05) differ from the hest-performing method on each task.
We train a policy x (s), (g))), in which we propagate the behavior cloning loss throughout the entire network. Both the state and goal encoders are MLPs with identical architecture. We detail the configuration in 5. This is to simulate the ResNet architecture and CLIP embeddings we use from real-world policy training. We define separate state and goal encoder ((s) and vigg,, and we modify LTRA as:
CIRA LBc((ss, 0z st )i= o. v) + a s
1i=1: 2/)
(21)
The rest of the implementation are carried over from OGBench. We evaluate each method with 10 seeds, and we take the final evaluation epoch per seed calculate the average success rate, the same way OGBench calculates success rate for its baselines. While we used a = in real world experiments, consistent with implementation from [27]. we adjust our a value in OGBench, it is hyperparameter. We report our optimal configuration in Table 5.
Note that = 0 turns the formulation into a version of GCBC with different architecture; denote this GCBC-Ø. We compare the performance of GCBC and GCBC-6 here across the 7 environments using table 4. Although the second formulation is parameterized than the original GCBC configuration, they have similar performances across the environments that we have evaluated on the performance of TRA does not rely on extra parameterization, but learning structured temporal representation.
We report the value of hyperparameters in table 5. The rest of the relevant hyperparameters are implemented from OGBench unless specified in the table.
H Statistical Analysis
Statistical significance in Table 1 computed with one-sided Welch's t-test, comparing the success rates of the highest scoring method with the other methods for each task. Highlighting was then determined using a p-value of 0.05. The p- and t- values for the language conditioned evaluations in Table are included in Table 6, and those for the goal conditioned evaluations in Table 7.
We also report the aggregated statistical comparisons of language-conditioned and image-conditioned performance across task sets in Table 8 and Table 9, respectively.
20

---

Table 7: Statistical comparisons of image-conditioned success rates
t-values
p-values
Task
Best Score
TRA
GRIF
AWR GCBC
Octo
TRA
GRIF
AWR
GCBC
Octo
(A)
open the drower
0.80
0.77
0.63
0.00
1.26
1.15
0.229
0.273
0.500
0.121
0.139
mushroom drower
0.90
0.00
1.89
1.13
0.45
0.00
0.500
0.057
0.152
0.335
0.500
(A)
close crower
1.00
0.00
2.45
1.00
2.45
2.45
0.500
0.035
0.187
0.035
0.018
take the item the crower
0.40
0.00
2.45
2.45
2.45
0.77
0.500
0.018
0.018
0.018
0.229
pul the apoonn on fowels
1.00
0.00
4.00
1.63
1.63
3.67
0.500
0.008
0.089
0.089
0.003
(B) put the spoons on the plates
0.90
0.00
9.00
0,45
1.89
9.00
0.500
<0.001
0.335
0.057
<0.001
pul the pom and sushi on plate
0.70
0.00
4.58
1.85
1.99
4.58
0.500
0.001
0.040
0.040
0.001
(C)
suhi and mushroom bowl
0.60
0.43
3.67
0.00
1.55
0.68
0.337
0.003
0.500
0.077
0.258
(C) com, banona, and sush in bowl
0.50
0.00
3.00
0.00
3.00
0.34
0.500
0.007
0.500
0.007
0.372
(D)
com on plate then sushi in pot
0.30
0.00
0.40
1.96
1.96
1.96
0.500
0.350
0.041
0.041
0.041
sweep right
0.80
0.40
1.26
4.00
4.00
0.00
0.350
0.121
0.008
0.008
0.500
(B) told cloth into the center
0.80
0.00
6,00
6.00
6.00
0.72
0.500
<0.001
<0.001
<0.001
0.249
move ball popper and swoop towal
0.60
0.00
1,55
3.67
1.55
0.68
0.500
0.077
0.003
0.077
0.258
each task the reported I-values compare the candidate method to the pext-pero method. We bold the t/p-values corresponding to methods that do not significantly (p 2 0.05) differ from the hest-performing method each task.
Table 8: Statistical comparisons of language-conditioned performance aggregated across task sets
t-values
p-values
Task Set
Best Score
TRA LCBC
GRIF
Octo
AWR
TRA
LCBC
GRIF
Octo
AWR
One-Step
0.77
0.00
2.38
2.38
2.78
3.22
0.500 0,0229 0.0229 0.0087 0.0028
Task Concatenation
0.80
0.00
5.36
5.36
5.36
6.25
0,500 0.0000 0.0000 0.0000 0.0000
Semantic Generalization
0.83
0.00
3.48
-12.04
7.98
3.47
0.500 0.0021 0.0000 0.0000 0.0010
Dependency
0.65
0.00
2.61
5.94
1.27
3.71
0.500 0.0165 0.0000 0.2203 0.0010
For each set of tasks indicated, report the t-values and p-values comparing the candidate method to the best- performing method, aggregated across all tasks in the set. We bold the t/p-values corresponding methods that do not significantly (p 2 0.05) differ from the best-performing method each task set.
Table 9: Statistical comparisons of image-conditioned performance aggregated across task sets
t-values
p-values
Task Set
Best Score
TRA
AWR GCBC
GRIF Octo|
TRA
AWR
GCBC
GRIF
Octo
One-Step
0.77
0.00
2.38
2.00
2.38
0.80
0.500
0.0115
0.0265 0.01 0.2137
Task Concatenation
0.82
0.00
3.79
4.32
7.89
4.03
0,500
0.0003 0.0001 0.0000 0.0001
Semantic Generalization
0.57
0.00
0.77
3.35
6.16
2.00
0.500
0.2234 0.0009 0.0000 0.0270
Dependency
0.35
0.00
3.20
3.20
1.69
1.69
0.500 0.0024 0.0024 0.0519 0.0519
For each set of tasks indicated, report the t-values and p-values comparing the candidate method to the best- performing method, aggregated across all tasks in the set. We bold the tip-values corresponding to methods that do not significantly (p 2 0.05) differ from best best-performing method on each task set.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, provide the theoretical and experimental results to show that we can learn task representations that enable compositionality.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, discuss the limitations of our work in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
22

---

Theoretical claims and assumptions are in Section 3.5, and proofs are in Appendix
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Experimental details are in Section 4, and the code is linked in Appendix A.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
23

---

Answer: [Yes]
Justification: Yes, the code is linked in Appendix A, and experimental details are in Appen- dices and C.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.c ec/
ssssonoliccyy Policy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /nips. cc/ wcc//ublic/ggee wetggddeeeoeeeuubbeeeeeeee iemiissionnl asionPolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experimental details are in Section 4, and the code is linked in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Significance is marked in figures and tables, with details in Appendix H.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
This is provided in Appendix B.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https ttttps//neeus http///eerrppppeeeee http:////cc/eiiiie....eeeeeeeeee
Answer: [Yes]
Justification: Yes, all research conducted in this paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, in Section 5 discuss the societal implications of future work based on TRA.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
25

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: We discuss safety in the context of future work, but the actual models used in this paper are pose no feasible risk.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, see citations and licenses in Appendix A.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thoode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
26

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: See Appendix A.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects were involved in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
No human subjects were involved in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
27

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: Pre-trained models used for data augmentation and representations are de- scribed in the paper.
Guidelines:
The answer NA means that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: : /ttp///eissp///eeee /ttpa//ips.ce/ooeeeeeeeeeeeeeeee /ttp:////cceeeennee//////111 for what should or should be described.
28


Reviews:
[{"review_id": "FeBmdgikKq", "reviewer": "Reviewer_nx61", "summary": "This paper introduces Temporal Representation Alignment (TRA), an auxiliary contrastive objective that aligns representations of current states, future goals, and language instructions in a shared embedding space. By adding symmetric NCE losses on state–future‐state pairs and goal–language pairs to a standard behavior‐cloning policy, TRA enables zero‐shot “stitching” of known subtasks into novel, long‐horizon manipulation sequences. The authors demonstrate over 40 % relative gains on compositionally‐out‐of‐distribution tasks in a real WidowX250 setup (5 Hz execution) and show consistent improvements across seven environments in the OGBench benchmark—outperforming both imitation‐only and offline‐RL baselines on stitching tasks . A theoretical bound (Theorem 1) further relates in‐distribution to compositional error under successor‐feature assumptions.", "questions": "1. How does varying the geometric sampling parameter γ affect compositional success? \n\n2. Can you provide t-SNE or clustering metrics showing that subtasks and their compositions occupy distinct, composable regions in the learned latent space?\n\n3. Have you evaluated TRA on a different manipulator (e.g., a Panda arm) or at higher control rates?\n\n4. How does TRA handle paraphrased or out-of-distribution instructions?", "limitations": "Please see the weaknesses.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. Minimal, Orthogonal Extension: TRA simply augments a standard goal- and language-conditioned BC policy with two lightweight contrastive losses, requiring no additional planner or value network.\n\n2. Strong Compositional Gains: Real-world experiments show a > 40 % boost on out-of-distribution “stitch” tasks, and OGBench results confirm superior performance on environments demanding composition.\n\n3. Unified Language + Vision: The same auxiliary losses work for both goal-image and natural-language instructions, demonstrating modality-agnostic compositionality.\n\n4. Theoretical Backing: Theorem 1 and Corollary 1.1 provide clear guarantees that temporal alignment reduces compositional generalization error under mild assumptions.\n\nWeaknesses:\n\n1. Limited Physical-Robot Diversity: All hardware tests use a single 5 Hz WidowX250 arm. It remains unclear whether TRA scales to faster control loops (e.g., 20 Hz Panda) or different kinematics.\n\n2. Some lacking ablation: Key choices—geometric horizon γ in the future‐state sampling, temperature and weighting of the contrastive losses—lack ablation.\n\n3. Failure Modes Under Multimodality: As noted, TRA struggles with inherently multimodal tasks (e.g., when multiple valid grasps exist), but no mitigation strategies are proposed .", "was_revised": true, "final_justification": "Thanks for the author's rebuttal. I don't have additional questions. I will keep my score.", "rebuttal": "Thank you for your thoughtful response. We will answer your questions below. **Please let us know if you have any additional questions or concerns.**\n\n> How does varying the geometric sampling parameter γ affect compositional success?\n\nWe have run an additional ablation of $\\gamma$ in OGBench in the `humanoidmaze-medium-stitch` environment, fixing the alignment coefficient at 40.\n\n| $\\\\gamma$      |   Mean (%) | Std. Err (%) |\n| :-----------: | ---------: | -----------: |\n| 0.8  |       28.3 |         2.2 |\n| 0.9  |       35.5 |         2.4 |\n| 0.95  |       42.8 |         2.6 |\n| 0.99  |       46.1 |         1.9 |\n| 0.995  |       41.8 |         2.2 |\n\nValues of $\\\\gamma$ between $\\\\approx 0.95$ to $0.995$ seem to perform well, which is similar to CRL's sampling coefficient of 0.995 [1].\n\n> Can you provide t-SNE or clustering metrics showing that subtasks and their compositions occupy distinct, composable regions in the learned latent space?\n\nThank you for this suggestion. We cannot include figures in this response, but we will add this t-SNE visualization to our revised manuscript.\n\n> Have you evaluated TRA on a different manipulator (e.g., a Panda arm) or at higher control rates?\n\nBridgeData only contains one embodiment for training, which is the WidowX setup we have in our experiment. We will revise our future work section to discuss additional embodiments and datasets, such as DROID [2] or OXE [3].\n\n> How does TRA handle paraphrased or out-of-distribution instructions?\n\nWe have done additional experiments on the task ''put the sushi, corn, and the banana in the bowl'', in which we have tried both ''put everything in the bowl'' and ''put the sushi, corn, and the banana in the bowl'' as our instructions. ''Put everything in the bowl'' has a success rate of 8 out of 10 tries, and ''put the sushi, corn, and the banana in the bowl'' has a success rate of 6 out of 10 tries. \n\nIn practice, using an effective pretrained language encoder, such as the CLIP model we use, is the component that ensures robustness to these small perturbations in the instruction format.\n\n---\n\n[1] Park, S. et al., 2025, ''OGBench: Benchmarking Offline Goal-Conditioned RL''. ICLR.\n\n[2] Khazatsky, A. et al., 2024, ''DROID: A Large-Scale In-the-Wild Robot Manipulation Dataset''. RSS\n\n[3] Open-X Embodiment Collaboration, 2024, ''Open-X Embodiment: Robotic Learning Datasets and RT-X Models''. ICRA"}, {"review_id": "wfyUidLgQm", "reviewer": "Reviewer_3j6m", "summary": "This paper introduces Temporal Representation Alignment (TRA), a novel approach for learning useful representations for solving goal/instruction-directed robotics tasks. The fundamental problem addressed by this paper is how to allow agents learning from a dataset of demonstrations to compose solutions to in-distribution tasks to solve longer-horizon out-of-distribution tasks. For example, if the dataset contains solutions to the tasks $s \\rightarrow w$ and $w \\rightarrow g$, how should an agent learn a representation that enables it to compose these solutions to solve the task $s \\rightarrow g$?\n\nThe proposed approach treats this purely as a representation-learning problem, and solves it by ensuring that the representations of states $s$, $w$, and $g$ are aligned according to the temporal structure of the demonstration trajectories. By mapping states, goals, and natural-language instructions into a shared embedding space that captures this temporal alignment, the agent can naturally “stitch together” demonstrated behaviours to perform composite tasks.\n\nThe authors demonstrate the performance of TRA on a diverse set of long-horizon robotics tasks with positive results, particularly in cases where solving tasks requires the sort of behaviour composition that TRA is well-suited to.", "questions": "I found the discussion around Equation 4 – of how the proposed approach allows an agent to compose information from demonstration trajectories $s \\rightarrow \\ldots \\rightarrow w$ and $w \\rightarrow \\ldots \\rightarrow g$ to capture the higher-level behaviour $s \\rightarrow \\ldots \\rightarrow g$ – to be particularly useful. This kind of sequential composition is a key idea motivating many existing hierarchical reinforcement learning methods. For instance, HIRO [1] learns a high-level policy that proposes a sequence of subgoals that a low-level policy navigates between, effectively stitching shorter segments together to solve a higher-level task.\n\n**How does the proposed approach compare with hierarchical methods (like HIRO) that attempt to address this problem of sequentially composing (or, “stitching together”) behaviours?**\n\nPlease note that I raise this not to diminish this paper’s contributions; on the contrary, I believe such discussion would further highlight the originality and significance of the proposed approach and better set it apart from related work in hierarchical reinforcement learning.\n\n---\n\nWhile the authors target robotics problems with continuous state and action spaces, I suspect many of the ideas proposed in this paper might be more generally applicable.\n\n**Do the authors foresee any challenges when adapting the proposed approach to work in domains with discrete state or action spaces?**\n**If so, do they have any initial ideas about how these issues might be addressed?**\n\n---\n\nThe reported success rates compare models trained on identical amounts of data, but it would be useful to know how the learned representations change as the amount of training data increases. If the proposed approach can learn useful representations with less (or more, as the case may be) data than alternative approaches, that would be useful to know.\n\n**Have the authors investigated the sample-efficiency of the proposed method?**\n\n---\n\n**References**\n\n[1] Nachum, O., Gu, S.S., Lee, H. and Levine, S., 2018. Data-efficient hierarchical reinforcement learning. Advances in neural information processing systems, 31.", "limitations": "The authors are candid about the limitations of the proposed approach, including a thoughtful discussion of these limitations and highlighting promising avenues for future work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Originality & Significance**\n\nThis paper addresses the important problem of allowing robotic agents to compose behaviours from large demonstration datasets to solve complex, long-horizon tasks that exceed the temporal scope of any single trajectory in their training data. By tackling this “trajectory stitching” problem directly and treating it as part of the representation learning problem, the proposed approach enables performance improvements over existing approaches across a wide range of challenging long-horizon tasks.\n\nThe proposed approach is also positioned well for widespread adoption. It is conceptually simple, seems straightforward to implement, and requires no additional information beyond what is already included in common behaviour-cloning datasets. Furthermore, its applicability is broadened by the fact that it can condition behaviour based on either raw goals or natural-language instructions.\n\nOverall, I believe the ideas presented in this paper represent a significant step forward for goal-directed learning from demonstrations – and perhaps representation learning for sequential decision problems more broadly – and should prove a valuable tool for the community.\n\n---\n\n**Quality**\n\nThe quality of this submission is very high.\n\nThe evaluation is both thorough and focused. Centred around four specific research questions, it demonstrates the performance of the proposed approach across a diverse set of challenging real-world tasks, and includes comparisons to a wide range of relevant baselines.\n\nIt was particularly refreshing to see the honesty with which the authors presented their results. For instance, Table 1 highlights the best-performing methods *up to statistical significance*, rather than simply highlighting the absolute highest performance (which, in many cases, would have made the proposed approach look even better), and Section 4.5 explicitly discusses the proposed approach’s failure modes.\n\nAlongside the thorough empirical evaluation, the method is also well-principled and supported by convincing theoretical work. I appreciated the care the authors took to introduce, motivate, and explain each element of their representation-learning objective.\n\n---\n\n**Clarity**\n\nFor the most part, the paper is written very clearly.\n\nThe authors explained the intuition behind the proposed approach well, and made good use of examples to help motivate it. The paper contains enough detail for me to understand the proposed method well. Taken together with the implementation details included in the appendices, I am confident that the authors have provided enough information to allow an expert reader to reproduce their results.\n\nThere are, however, a couple of potential issues with mathematical notation, which I have discussed in detail in the “Paper Formatting” section.", "was_revised": true, "final_justification": "Following discussion with the authors, I am happy to maintain my current rating for this paper and advocate for its acceptance.\n\nI already had a favourable opinion of this paper: the proposed approach represents a step forward in representation learning for goal-directed tasks, and the evidence presented for its effectiveness is compelling.\n\nDuring the discussion, the authors clarified additional properties of their method (e.g., natural extensibility to domains with discrete action spaces and favourable sample complexity), discussed additional experiments and comparisons (e.g., to hierarchical baselines), and addressed minor formatting and notational issues.\n\nLike other reviewers, I initially questioned why no comparisons to hierarchical methods were included in the paper; I am happy to hear that the authors are going to include such comparisons in the camera-ready version. However, unlike some other reviewers, I believe the contributions and supporting evidence are already strong enough to merit acceptance without these additional results, which I view as a welcome addition rather than a hard requirement.\n\nI also wish to note that some other reviewers commented on the proposed method not always significantly outperforming existing approaches across all evaluation domains. I do not support this criticism: it is not healthy for our community to dismiss novel ideas just because they are not dominant everywhere. In fact, I commend the rigour with which the authors performed their experiments and the candour with which they presented their results and limitations.\n\nI thank the authors for their high-quality submission and their engagement during the rebuttal and discussion period.", "rebuttal": "Thank you for your thoughtful response. We will address the questions you have raised below. **Please let us know if you have any additional questions or concerns.**\n\n> How does the proposed approach compare with hierarchical methods (like HIRO) that attempt to address this problem of sequentially composing (or, ''stitching together'') behaviours?\n\nHierarchical methods such as HIRO [1] often propose subgoals as high-level actions, represented as states to reach or, to instruct a low-level policy to reach goals. While hierarchical methods can yield impressive results, they introduce ''moving parts'' that can fail, and often require additional assumptions such as access to strong pretrained models [2,3]\n\n> Do the authors foresee any challenges when adapting the proposed approach to work in domains with discrete state or action spaces? If so, do they have any initial ideas about how these issues might be addressed?\n\nThe method can be extended naturally to discrete environments. For discrete action spaces, a cross entropy loss could be used instead of MSE in the actor loss. Such an approach would also enable the use of autoregressive transformer-based architectures for the actor, that could be more effective in environments with multimodal behaviors. Likewise, discrete state spaces could be handled with simpler MLP architectures or through architectures like transformers.\n\n> Have the authors investigated the sample-efficiency of the proposed method?\n\nIn BridgeData, we have trained TRA using 150,000 steps of training with a batch size of 128. Given that we sample an additional future state using geometrically for each data point, in the end we use a similar amount of data as the GCBC/LCBC baselines, which train for 300,000 steps with the same batch size. Octo [4] on the other hand, uses the entirety of OXE [5], which is much more data-intensive than the rest of the approaches.\n\nOn OGBench, we use the same data loading procedure as CRL, which means that it requires double the amount of data required for GCBC, and the same amount of data as CRL.\n\n---\n\n[1] Nachum, O. et al., 2018. ''Data-Efficient Hierarchical Reinforcement Learning''. NIPS.\n\n[2] Belkhale, S. et al., 2024. ''RT-H: Action Hierarchies Using Language''. \n\n[3] Ahn, M. et al., 2022. ''Do As I Can, Not As I Say: Grounding Language in Robotic Affordances''. CoRL.\n\n[4] Octo Model Team, 2024, ''Octo: an Open-Source Generalist Robot Policy''. RSS.\n\n[5] Open-X Embodiment Collaboration, 2024, ''Open-X Embodiment: Robotic Learning Datasets and RT-X Models''. ICRA."}, {"review_id": "gSrXZpcRaC", "reviewer": "Reviewer_6ZRh", "summary": "This paper addresses the problem of compositional generalization in robotic skills, where elementary skills learned during training can be recombined to solve longer, more complex tasks at test time. The authors propose Temporal Representation Alignment (TRA), a method that achieves compositionality purely through learning aligned state representations via temporal alignment of state and goal/instruction embeddings, which are then used in a goal-conditioned policy. To this end, a contrastive loss is introduced to align representations from the same expert demonstration. The policy is trained via behavioral cloning on expert demonstrations using these temporally aligned goal embeddings. TRA is evaluated in both simulation and real-world robotics settings, and the results show improved performance compared to various goal- or language-conditioned baselines, particularly in compositional generalization to novel task sequences.", "questions": "1. Compositionality-specific baselines: None of the baselines used are designed to explicitly address compositional generalization. Could the authors include comparisons to hierarchical methods or planning-based approaches, as discussed in the related work?\n2. Clarification on rarely-seen skills: In the paragraph titled \"Does TRA help rarely-seen skills within the dataset?\", the actual frequency of the skills is not mentioned. How does the analysis support the claim about rarely-seen skills?\n3. Task set clarifications: Are all tasks from set A in-distribution training tasks? What is the difference between task set B and C? Why is the use of different objects with the same type referred to as \"semantic generalization\"?\n\nMinor Comments\n\n- Table formatting errors: In Table 1, some rows highlight 0.00 success rates (e.g., row 4 for AWR, row 6 for Octo) as statistically significant best results. These seem to be mistakes, please double-check.\n- Table references: Tables are referred to only by number in several places, which is confusing. For example, l. 211: \"1 shows\" → should be \"Table 1 shows\"\"; l. 611: \"in 5\" → should be \"in Table 5\".", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- addresses an important and challenging problem of compositional generalization in a novel way by introducing a temporal alignment loss over representations.\n- empirical evaluation: includes both simulation and real-world experiments, and demonstrates statistically significant performance improvements over a variety of baselines (even though the baselines don’t directly target compositionality!).\n- the paper is mostly well-written, and uses figures, algorithm boxes, and tables effectively to illustrate key concepts.\n\nWeaknesses:\n- some parts of the method description would benefit from more natural language explanation. For example, Eqs. 6 and 7 are presented without much intuition, and the roles of $\\psi$ and $\\phi$ (introduced in Eq. 5) are not clearly described in text until much later (Subsection 3.4).\n- no baseline  is directly targeting compositional generalization. While the method is compared to strong goal- and language-conditioned policies, it is not evaluated against approaches that specifically tackle compositionality, such as hierarchical policies or planning-based methods, e.g. the methods discussed in the related work section.", "was_revised": true, "final_justification": "The authors have appropriately addressed my questions regarding the dataset.\nHowever, my main concern regarding baselines remains: while the paper compares against several strong methods, none are specifically designed to tackle compositional generalization. It would have been helpful to understand how TRA fares against such approaches.\nThat said, I find the paper to be a meaningful and promising contribution, and I believe its strengths outweigh these concerns.", "rebuttal": "Thank you for your thoughtful feedback. We will address your questions and concerns in the section below. **Please let us know if you have any remaining questions or concerns.**\n\n> Clarification on rarely-seen skills: In the paragraph titled ''Does TRA help rarely-seen skills within the dataset?'', the actual frequency of the skills is not mentioned. How does the analysis support the claim about rarely-seen skills?\n\nThe following table shows the frequency of the instruction labels in the Bridge dataset that are paraphrases of the commanded evaluation tasks (out of 53,191 instructions).\n\n| Scene | Command                                                                                              | Count | Percentage |\n| :---: | :--------------------------------------------------------------------------------------------------- | ----: | ---------: |\n| A     | open the drawer                                                                                      |   500 |     0.940% |\n| A     | take the mushroom out of the drawer                                                                  |     1 |     0.002% |\n| A     | close the drawer                                                                                     |   435 |    0.818% |\n| B     | put spoon on plate                                                                                   |    26 |    0.049% |\n| B     | put the spoon on the towel                                                                           |   155 |    0.291% |\n| B     | fold the cloth into the center                                                                       |    42 |    0.079% |\n| B     | sweep the towels to the right of the table                                                           |    72 |    0.135% |\n| C     | put the food items on the plate                                                                      |    41 |    0.077% |\n| C     | put the food items in the bowl                                                                       |    41 |    0.077% |\n| C     | put everything in the bowl                                                                           |     1 |    0.002% |\n| D     | open the drawer and then take the mushroom out of the drawer                                          |     0 |    0.000% |\n| D     | move the bell pepper to the bottom right corner of the table, and then sweep the towel to the top right corner of the table |     0 |    0.000% |\n| D     | put the corn on the plate and then put the sushi in the pot                                          |     0 |    0.000% |\n\nNone of the tasks are more than 1% of the dataset, and for some tasks, only a few dozen related demonstrations exist out of almost 70,000 trajectories [1] (the discrepancy between the number of trajectories and instructions is from some of the dataset being unlabeled). While some similar behaviors to tasks in scenes A, B, and C appear in the dataset, none of the composed behaviors in scene D occur in the data.\n\nWe will revise the paper to include this table and the discussion above.\n\n> Task set clarifications: Are all tasks from set A in-distribution training tasks? What is the difference between task set B and C? Why is the use of different objects with the same type referred to as ''semantic generalization''?\n\nFor set A, both the drawer tasks are in the support of the dataset (see table above), though the actual evaluation scene features objects in different positions from those seen in the dataset. Only one trajectory that takes the mushroom out of the drawer is in the dataset.\n\nThe main difference between set B and set C is that set B tasks involve performing the same subtask in sequence, but in set C, semantically distinct subtasks must be executed in sequence. For instance, the agent only needs to put multiple spoons on multiple towels in set B, but the agent needs to put multiple different kinds of food items in set C. Set C is more challenging than B because it depends on successful execution of distinct subtasks, and further requires subtasks to be executed in the correct order for success.\n\n> no baseline is directly targeting compositional generalization. While the method is compared to strong goal- and language-conditioned policies, it is not evaluated against approaches that specifically tackle compositionality, such as hierarchical policies or planning-based methods, e.g. the methods discussed in the related work section.\n\nWe chose to focus this work on compositional generalization in non-hierarchical models, and thus excluded hierarchical and planning-based approaches. While often effective, these approaches introduce separate challenges, such as needing a much larger policy network [2], a small set of demonstrations during test-time [3], or access to powerful vision-language models [3,4]. Our approach highlights that compositionality can be achieved without using explicit planning or decomposition.\n\n> Table formatting errors\n\n> Table references: Tables are referred to only by number in several places ...\n\nThank you for noting these errors, we have fixed them in the revision.\n\n---\n\n[1] Walke, H. et al., 2023. BridgeData V2: A Dataset for Robot Learning at Scale. CoRL\n\n[2] Belkhale, S. et al., 2024. ''RT-H: Action Hierarchies Using Language''. \n\n[3] Myers, V. et al., 2024. ''Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation''. CoRL.\n\n[4] Ahn, M. et al., 2022. ''Do As I Can, Not As I Say: Grounding Language in Robotic Affordances''. CoRL."}, {"review_id": "HdNZVVuNRj", "reviewer": "Reviewer_Liiq", "summary": "This paper proposes temporal representation alignment (TRA), a training approach for a language-conditioned and goal-conditioned policy which uses contrastive losses to align future (image-based) goal representations with current (image-based) state representations and text-based task specification representations. Intuitively, this enables the policy to be trained on two separate trajectories where Trajectory A’s goal state is the initial state of Trajectory B, then be able to compose these trajectories together conditioned only on Trajectory B’s goal state (or text instructions).\n\nThe paper draws from BridgeDataV2 to train TRA and several representative baseline policies for real-world execution. Evaluation task categories carefully separate tasks by their difficulty of compositional generalization, from simple one-step tasks to concatenated tasks, order-independent concatenated tasks with different objects of the same type (e.g., placing several food items into a container), and tasks with ordering dependencies between subtasks (e.g., opening a drawer then taking an object out of it). In these real-world experiments, the paper demonstrates that TRA achieves top success rates in all studied task categories in both language-conditioned and goal-conditioned settings. While policies from prior work are not statistically significantly worse than TRA in many cases, TRA achieves the best statistically significant performance in more categories than any other approach. Performance gains over the best baseline are most pronounced in tasks requiring compositional generalization. TRA vastly outperforms traditional offline reinforcement learning approaches (which are typically thought necessary to “stitch” subtasks together), and does not significantly benefit from being combined with such approaches. Together, these suggest that TRA is a remarkably effective approach despite its relative simplicity.\n\nThe paper goes on to evaluate TRA in 7 different simulated environments from OGBench. Consistent with the above results, TRA achieves the significantly best success rate in most environments (5 of 7), including 4 of 5 environments that require stitching together subtasks. The paper then briefly discusses TRA’s failure cases.", "questions": "Based on my Weakness 2, I ask one question to facilitate the rebuttal: **Do compositionally-OOD tasks truly consist of less common subtasks?** Given clarification on this in the rebuttal, it’s possible I will increase my quality score up to 4\\. \n\nMy other weakness regarding Figure 1 cannot really be addressed without seeing a new version of the Figure, which the authors are welcome to discuss or present. That said, Figure 1 did not significantly impact my score (clarity score is already 4).\n\nGiven that the results are mixed in some cases (i.e., don’t consistently significantly outperform baselines), I don’t anticipate raising my significance score beyond 3 (which is a positive score). Given that the proposed approach is fairly simple, I don’t anticipate raising my originality score beyond 3 (which is a positive score). For these reasons, I don’t anticipate raising my overall score, which is already quite positive.", "limitations": "The discussion of limitations of the approach in Section 5 is sufficiently thoughtful and provides areas for future work to improve upon. The discussion of failure cases in the main body of the paper is quite brief. I don't see much discussion about the broader society impact in the paper.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strength 1: Paper is well written and clear.** \n\nThis paper is well written. It clearly conveys its motivation, contributions, methods, experimental design, results, and conclusions. The paper clearly and intuitively contextualizes its contributions with prior work in language- and goal-based robotic manipulation, compositional generalization in sequential decision making problems, and representation learning for physical states and procedures. It is quite polished with limited typos, and its figures and notations are mostly very clear and illustrative. Extensive appendices provide lower-level supplementary details about the datasets, methods, implementation, and analyses for reproducibility. In terms of clarity, it is among the top tier of submissions.\n\n**Strength 2: Experiments soundly demonstrate effectiveness of TRA.**\n\nThe experimental results are motivated by clear and well-defined research questions. These questions are mostly effectively answered by experiments, demonstrating that TRA indeed improves zero-shot composition of tasks over prior methods, better captures rare skills in the dataset, and works even without explicit subtask planning or reinforcement learning. Performance improvements over prior work are sharp and significant. Experiments cover both language-conditioned and goal-conditioned execution, and simulated and physical settings. Evaluations carefully divide BridgeData according to difficulty in terms of compositional generalization. To the best of my knowledge, the baselines are sufficiently representative of approaches used in prior work. \n\n**Weakness 1: Figure 1 does not provide a clear impression of the task or method.**\n\nWhile this is a relatively minor concern, I found that Figure 1 left a lot of questions for me after reading Section 1 of the paper (where it’s referenced) for the first time. Specifically, I had the following questions (these do not require an answer, as I figured them out after reading the entire paper, but hopefully can help guide revisions of the figure):\n\n* What is the task that TRA, AWR, and BC are trying to perform? The caption claims this is a “language task”, but it clearly involves manipulation too.  \n* What are we trying to observe through the comparison with AWR and BC? What are AWR and BC doing wrong, and how does it relate to compositionality? I see that AWR only puts one food in the bowl, but does this show the entire execution of the robot? Further, it seems that BC is unable to find any objects. What I don’t understand from the figure and caption is how this demonstrates a lack of “compositionality.”  \n* What is the TRA method doing? How is the time-contrastive loss calculated, and what does that have to do with the sequence of frames? Providing an intuitive sentence or phrase in the caption to describe the method at the high level would help a lot.\n\nAlso, the acronym for advantage weighted regression (AWR) is only defined in the appendix \\- I would recommend defining this somewhere in the main body of the paper.\n\n**Weakness 2: Paper seems to conflate “rarely-seen” and “challenging” tasks.**\n\nOne key question for the experimental results was “How well does TRA capture skills that are less common within the dataset?” However, the discussion around this at L241-246 seems to focus on the more challenging tasks in terms of compositional generalization (from Sets C and D involving semantic generalization across objects and subtask dependencies), which are compositionally OOD. However, can we safely assume that subtasks from Sets C and D are less common in training than others? I can’t seem to find details to support this, and would recommend that the authors better clarify this in the revised version of the paper, and consider rephrasing this research question if needed.", "was_revised": true, "final_justification": "The concerns raised in my review have been addressed by the authors' response. However, given that the results are mixed in some cases (i.e., don’t consistently significantly outperform baselines), and given that the proposed approach is fairly simple, I will maintain my score as expected (although I do view it slightly more favorably now).", "rebuttal": "Thank you for your detailed and constructive response. We will address some of your concerns below.\n\n> Weakness 2: Paper seems to conflate ''rarely-seen'' and ''challenging'' tasks ... Based on my Weakness 2, I ask one question to facilitate the rebuttal: Do compositionally-OOD tasks truly consist of less common subtasks? Given clarification on this in the rebuttal, it's possible I will increase my quality score up to 4.\n\nIn practice, almost all rarely-seen tasks are challenging for policies trained on the Bridge v2 dataset [1]. The following table shows the frequency of instructions that are semantically similar to the evaluation tasks (same instruction label up to equivalent paraphrases) found in the dataset (of 53,191 instructions).\n\n| Scene | Command                                                                                              | Count | Percentage |\n|:-----:|:-----------------------------------------------------------------------------------------------------|------:|-----------:|\n| A     | open the drawer                                                                                      |   500 |    0.94% |\n| A     | take the mushroom out of the drawer                                                                  |     1 |    0.002% |\n| A     | close the drawer                                                                                     |   435 |    0.82% |\n| B     | put spoon on plate                                                                                   |    26 |    0.05% |\n| B     | put the spoon on the towel                                                                           |   155 |    0.291% |\n| B     | fold the cloth into the center                                                                       |    42 |    0.079% |\n| B     | sweep the towels to the right of the table                                                           |    72 |    0.135% |\n| C     | put the food items on the plate                                                                      |    41 |    0.077% |\n| C     | put the food items in the bowl                                                                       |    41 |    0.077% |\n| C     | put everything in the bowl                                                                           |     1 |    0.002% |\n| D     | open the drawer and then take the mushroom out of the drawer                                          |     0 |    0% |\n| D     | move the bell pepper to the bottom right corner of the table, and then sweep the towel to the top right corner of the table |     0 |    0% |\n| D     | put the corn on the plate and then put the sushi in the pot                                          |     0 |    0% |\n\nNone of the tasks comprise more than 1% of the dataset. Note that while some demonstrations corresponding to tasks in scenes A, B, and C do appear in the Bridge data, the complex behaviors in scene D do not; they are **compositionally-OOD**. **Does this clarification address your concern?**\n\n> Weakness 1: Figure 1 does not provide a clear impression of the task or method. My other weakness regarding Figure 1 cannot really be addressed without seeing a new version of the Figure, which the authors are welcome to discuss or present. That said, Figure 1 did not significantly impact my score (clarity score is already 4).\n\nThank you for your feedback. We will modify the caption to better define terms (language task, compositionality), and expand the TRA portion of the figure to show how frames are being aligned across time / with language. Please let us know if there are any other changes to the figure that might be helpful.\n\n> Missing references, typos, definitions in paper formatting concerns.\n\nThank you for noticing these issues. We will correct them in our revision.\n\n---\n\n[1] Walke, H. et al., 2023. BridgeData V2: A Dataset for Robot Learning at Scale. CoRL"}]

Author rebuttal:

Meta-review: In this paper, the authors propose Temporal Representation Alignment (TRA), a contrastive objective for robot learning that aligns state, goal, and instruction representations within a shared embedding space. The method leverages symmetric NCE losses between current–future states and between goals–language to a behavior-cloning policy, aiming to enable “stitching” of demonstrated subtasks into novel long-horizon tasks. Experiments in the paper show relative improvements on out of distribution tasks in real-world robot manipulation and consistent gains on benchmarks such as OGBench; particularly shining in tasks that require compositional generalization.

Initial reviews were positive. Strengths included the simplicity of the extension (adding two lightweight losses to normal behavior cloning), gains on compositional generalization on both real robots and simulation (R-nx61, R-3j6m), being applicable to both goal images and natural language (R-nx61), and theoretical guarantees (R-nx61). Weaknesses raised were the limited hardware diversity (all results being on a single arm running at 5 Hz, hence questions on whether it can scale to higher frequencies), missing ablations for hyperparameters such as the geometric horizon and loss weights, and observed failure modes on multimodal tasks.

The rebuttal provided ablations on geometric horizon, showing stable performance for different values, added t-SNE visualizations of representation clustering, and reported paraphrase robustness tests. Authors clarified sample efficiency (comparable or better than baselines at equal training steps) and argued that TRA naturally extends to discrete action/state settings. They acknowledged hardware limitations due to the choice of datasets used, and committed to discussing generalization to other embodiments in a future version. 
Reviewers continued to be positive across the board post-rebuttal: R-nx61 maintained an “accept” recommendation, and R-3j6m has advocated for acceptance after noting clarifications on hierarchical comparisons, discrete extensions, and efficiency.

Overall, TRA introduces a simple auxiliary objective that improves compositional generalization without requiring planning or hierarchical structure. The paper is supported by both theoretical analysis and experimental results across simulation as wel as real world. The reviewers were aligned in support after rebuttal. I do recommend acceptance, but also reaffirm that adding comparisons with hierarchical baselines such as HIRO / HIQL would be a really useful addition to the camera-ready version.
Decision: Accept (poster)