Zero-Shot Policy Transfer with Disentangled Task Representation of Meta-Reinforcement Learning — part 5/6

Cleaned-text code points [58325, 72074). Bibliography omitted here; source.md includes it.

Under review for the Reinforcement Learning Conference (RLC)
agent crosses a true-cell boundary the cursor teleports to an unrelated observation cell (Fig-
ure 10). A continuous trajectory in the true grid thus appears to the agent as a sequence of
action-conditioned jumps, and the agent must implicitly invert the cell permutation to navigate.
Figure 10: An example
ColorGrid transition for
n= 5.
Observations.The agent receives a three-dimensional position obser-
vation. Its first two entries are the permuted position cursor, formed by
relabelling the integer part of the true position bycr and appending the
unchanged sub-cell fraction. The third entry is a soft-reset flag that fires
for one step at the boundary between inner episodes within a trial. Both
cd andc r are hidden.
Task distribution.Permutations and reward grids are drawn from two
deterministic deduplicated pools, with reward entries sampled iid from
N (0, 3)and clipped to[ −10, 10]. Each pool is partitioned into10,000
training entries and a disjoint10,000held-out test entries. Thei-th task
is( ci
d,ci
r), drawing one entry from each pool in lockstep. In Figure 1,
tasks are sampled from the cross product of the two pools, so that policies
can be trained across many reward functions in a fixed permutation or
vice versa.
Trial structure and metric.A meta-trial bundlesk inner episodes
that share the same(cr,cd). Inner episodes have a maximum length of60
steps forn = 5and80steps for n = 9. Between inner episodes the agent’s
position is re-randomized, with the soft-reset flag firing for one step. The
trial truncates after thek-th episode. We report oracle-normalized return,
defined as per-episode return divided by that of an oracle which teleports
from the random start cell toarg maxcr and remains there. A score of
1.0matches the oracle.
B.1.2 Wind
Wind is a 2D point meta-RL navigation task introduced by Ni et al. (Ni et al., 2022) as part of their
POMDP suite. The agent commands position deltas on the plane, and the world adds a hidden,
task-specific wind vector to its motion at every step. A task is an unobserved(w,g )pair in which w
perturbs the dynamics whileg defines the reward landscape. The two hidden variables stress meta-RL
along complementary axes. The wind is identifiable only from the difference between commanded
actions and observed motion, and the goal only from reward, so a competent policy must jointly
infer the dynamics and search for an unmarked target.
Dynamics and reward.The agent’s state is a 2D position st∈R 2, advanced by an action
at∈ [−0.1, 0.1]2 and a per-task windw∈ [−0.08, 0.08]2 via st+1 = st +at +w. Inner episodes
last75steps and terminate early on entry into a small disk around the goalg. The reward is
rt =e−∥st−g∥ +1[st∈goal disk], bounded everywhere on the plane with a sharp+1bonus inside
the disk.
Observations.The agent receives its own 2D position together with an episode-boundary flag that
fires for one step at the boundary between inner episodes within a trial. Bothwandgare hidden.
Task distribution.During training, we sample80winds uniformly from[ −0.08, 0.08]2 and80
goal locations on the unit circle, with the circle partitioned into the four standard quadrants. One
quadrant (Q2 in our experiments) is held out from the training goal set, so the agent is never rewarded
for entering the upper-left region of the plane during training. The two test environments measure
distinct generalization gaps that we report separately.
17

Under review for the Reinforcement Learning Conference (RLC)
Trial structure and metric.A meta-trial bundlesk inner episodes that share the same(w,g ).
The agent’s position resets to the origin between inner episodes, with the soft-reset flag firing for
one step, but the task persists across the trial so the policy can use early episodes to identify(w,g )
and the rest to exploit them. We train withk = 3and evaluate atk = 35to probe how the policy
continues adapting beyond its training horizon. We report mean episodic return.
B.2 Metaworld
We base our experiments on the MetaWorld manipulation suite (Yu et al., 2020), with the reward
and goal-masking corrections of McLean et al. (2025), run through the GPU-accelerated MTBench
reimplementation (Joshi et al., 2025) for batched simulation. Tasks are simulated in NVIDIA
IsaacGym (Makoviychuk et al., 2021) with a Franka Panda arm, and episodes are 250 steps.
Tasks and goals.MetaWorld consists of 50 single-arm tabletop manipulation task families (sliding,
picking, button-pressing, door opening, etc.), each parameterized by a goal vector specifying a target
site or object position drawn from a task-specific range. Each training task family exposes 100
sampled goal variations, giving ML45 a total of4,500train configurations across 45 families and 5
held-out test families. The goal vector is masked to zero in the agent’s observation at every timestep,
so the agent must infer both the task family and the specific goal from interaction alone.
Action and observation spaces.The action is a 4-D continuous vectora∈ [−1, 1]4. Three
components specify a Cartesian end-effector delta and the fourth controls the parallel gripper.
The observation is a 39-D continuous vector containing the end-effector position and gripper, the
position and quaternion of up to two task-relevant objects (zero-padded for tasks with fewer), the
previous-timestep object pose, and the goal slot (masked as described above).
Reward.At each timestep, the agent receives the sum of a dense shaping term and a sparse success
bonus. MetaWorld reward functions have changed several times: MetaWorld v2 placed rewards on a
more consistent scale across tasks to avoid indirectly leaking task identity, while MTBench reduced
the dense reward magnitude by.01× to reduce returns and improve multi-task RL. Our experiments
avoid this issue with categorical value function networks, so we restore the original dense reward
scale, as dense rewards are expected to improve OptiBFM online task inference.
B.3 Humanoid Locomotion
Following Li et al. (2025), we create environments in IsaacLab (Mittal et al., 2025) for humanoids
with locomotion tasks. Basically, the humanoids are evaluated to do motion tracking on diverse and
challenging motion datasets, e.g. LAFAN (Harvey et al., 2020).
Although RL methods on robotics usually include domain randomization to achieve robust sim-to-real
transfer, we intentionally try to create more dynamic and random environments, for example, different
wind and gravity. So apart from the domain randomization on mass, friction and so on, we randomly
sample the scale and direction of wind, and gravity scale as shown in the figure below. Dynamics
setting details are shown in Table. 12.
B.3.1 Representations Details
State and Observation.The privileged state st∈R 463 includes root height, full body pose,
body orientation, as well as linear and angular velocities. The observable state is defined as
ot ≜
(
qt−¯q,˙qt, 1
4ωroot
t ,gt
)
∈R 64, whereqt∈R 29 is the joint position normalized by the nominal pose
¯q, ˙qt∈R 29 is the joint velocity,ωroot
t ∈R 3 is the root angular velocity, andgt∈R 3 is the projected
gravity. To mitigate partial observability, we use an observable historyot,H ≜(o t−H,at−H,...,o t)∈
R93H+64,which concatenates past observations and actions over a horizon of lengthH. All observation
components (except root height in the privileged state) are normalized with respect to the humanoid’s
current facing direction and root position. During BFMs training, we assume access to a dataset of
18

Under review for the Reinforcement Learning Conference (RLC)
Figure 11:Wind/Gravity illustration.During online data collection, the wind and gravity strength in
each environment are different – the horizontal arrows indicate the direction and strength of the wind (red as
the strongest, white as the weakest), while the length of vertical arrows indicate the strength of gravity.
Domain Randomization
Parameter Range
COM Offset [m]U([−0.02,0.02])
Link MassU([0.95,1.05])
FrictionU([−0.5,1.25])
Default Joint Pos [m]U([−0.02,0.02])
Push Robots [m/s]U([0,0.5])
Wind Direction[Rad]U([−π,π])
Wind MagnitudeU([0.0,30.0])
Gravity Magnitude[g]U([0.8,1.2])
Additive Observation Noise
Observation Range
qt−¯qU([−0.01,0.01])
˙qt U([−0.5,0.5])
gravt U([−0.05,0.05])
˙ωroot
t /4U([−0.05,0.05])
Figure 12: Details in training environment.
unlabeled motion trajectoriesD={τ}, where each trajectoryτ= (o 1,s 1,...,o T,sT )contains both
observable and privileged states.
Action.Following PHC, we use a proportional-derivative (PD) controller at each DoF of the
humanoid, with the actionat specifying the PD target. With the target joint set asqd
t =at, the
torque applied at each joint isτi =kp◦ (at−qt)−kd◦˙qt. This is different from the residual action
representation used in prior motion imitation methods, where the action is added to the reference
pose:qd
t = ˆqt +at to speed up training. We do not use any external forces or meta-PD control. For
the 29-degree-of-freedom (DoF) humanoid we use – Unitree G1, the actiona∈R 29 contains the PD
controller targets for all DoFs.
B.3.2 Illustrations of BFM’s Failure on Humanoid Locomotion Tasks
BFMs trained with multi-dynamics data might suffer from taking actions from the wrong environment
dynamics. This can be shown directly via the visualizations in Fig. 13. For BFM-Multi, in the stage
of trying to get up from the ground, the robot acts as there exists wind from certain directions,
loses balance and finally tracks badly. As a comparison, BFMs with context learner can identify the
dynamics, so the states are close to the BFM-Single, which is trained solely on one type of dynamics.
19

Under review for the Reinforcement Learning Conference (RLC)
Reference Motions
BFM-Single
BFM-Multi
Ours
Figure 13:Qualitative Comparison on Humanoid Motion T racking.
Algorithm 1 Latent Context Learning
Input:trajectory batch τ = {τ (i)}B
i=1 with τ (i) = {(s0,a 0,...,s T )}; EMA rate ρ; networks
Eθ,Sθ,LDMθ,IDMθ,FDMθ; EMA target encoderE−.
Output:augmented trajectory¯τ={(¯s t,at,¯st+1)}t with¯st = (st,ht,˜ct).
// Augment eachτwith dynamics features
ht =E θ(st,at−2), h−
t =E−(st,at−2),˜ct =S θ(τ:t),¯st = (st,ht,˜ct).
// Self-supervised dynamics losses
ˆLLDM =− 1
N
∑
i,t cos
(
LDMθ(ht,at,˜ct), h−
t+1
)
ˆLIDM = 1
2N
∑
i,t
‖‖IDMθ(ht,ht+1,˜ct,˜ct+1)−a t
‖‖2
ˆLFDM = 1
2N
∑
i,t
‖‖FDMθ(st,at,sg(˜ct))−s t+1
‖‖2
//˜ct detached by default (diagnostic only)
ˆLREG =− 1
B(B−1)T
∑
i̸=j,t
[
1−cos
(
˜c(i)
t ,˜c(j)
t
)]
Take a gradient step onθ to minimize a weighted sum of (normalized)ˆLLDM,ˆLIDM,ˆLFDM,ˆLREG.
EMA target update:E−←(1−ρ)E − +ρE θ.
returnaugmented trajectory¯τ.
C Implementation Details
C.1 Latent Context Learning
Algorithm 1 summarizes the Latent Context Learning (LCL) procedure, which augments states with
features used for BFM training and intrinsic rewards. Per-domain choices are summarized in Table 2.
Notably, MetaWorld uses a feed-forward dynamics encoder, while ColorGrid and Wind use small
Transformer encoders. The optional cross-trajectory regularizer is disabled in Section 4, though
preliminary experiments in Figure 4 enabled it.
20

Under review for the Reinforcement Learning Conference (RLC)
Table 2: Latent Context Learning hyperparameters.
Hyperparameter ColorGrid Wind MetaWorld
Latent feature dim (dimh) 32 4 12
Inverse-dynamics feature dim 32 4 12
Inverse-dynamics MLP hidden width 400 128 300
Latent-dynamics predictor hidden width 512 128 300
Diagnostic probe hidden width 400 128 256
Sθ (dmodel, layers, heads) Tformer (512, 6, 8) Tformer (128, 2, 4) FF (128, 2, –)
Shared across domains: latent-dynamics and inverse-dynamics loss weights both 1; EMA target rate ρ=5×10−3;
sphere-normalized features; cross-trajectory regularizer weight 0; probe gradients detached.
Table 3: BFM (TD-JEPA) hyperparameters.
Hyperparameter ColorGrid Wind MetaWorld
ϕ-encoder output dim (dimϕ) 200 160 220
ψ-encoder output dim (dimψ) 50 40 80
ϕ-predictor hidden width 1024 800 1024
ψ-predictor hidden width 1024 800 1024
Successor-feature discount (γsf) 0.97 0.97 0.98
Reward multiplier (SF target / BLR) 10 10 1
ϕ–ψorthogonality weight 0.10 0.25 0.10
In reference to notation in Bagatella et al. (2025). Shared across domains: two-tower predictor with three layers (one
shared trunk plus two per-head tower layers).
C.2 BFM
While many BFM updates could in principle be used interchangeably, our results use TD-JEPA
(Bagatella et al., 2025), primarily because it shares the most surface area with off-policy actor–critic
RL. The two headsϕ (state representation) and ψ (task representation) each use a two-tower
predictor, sharing a small input trunk before splitting into per-head towers. We additionally apply a
cross-trajectory orthogonality regularizer on the predicted features. A scalar reward multiplier (ten
for ColorGrid and Wind, one for MetaWorld) is applied to the per-step reward inside both the SF
prediction target and the BLR likelihood at deployment, so it sets the natural magnitude of the
regression targetr≈⟨ψ (s′),z⟩ and therefore couples directly to the Thompson-sampling noise level
σdiscussed in Appendix C.3.
C.3 Online Task Inference
Following OptiBFM (Rupf et al., 2025), we estimate the latent task vectorz∗ online with Bayesian
linear regression on the forward featuresϕ, resampling at every timestep. Our experiments use
Thompson sampling rather than the optimistic alternative because it scales to the long meta-rollouts in
our evaluations, and preliminary experiments found little performance gap. The posterior is governed
by an observation-noise scaleσ and prior precisionλ; we foundσ to be the key hyperparameter,
scaling roughly linearly with the typical magnitude of the scaled per-step reward.
21