Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization
Authors: Matthew Landers
Taylor W. Killian
Thomas Hartvigsen
Afsaneh Doryab
Venue: ICLR 2026 Poster (2026)
Keywords: reinforcement learning
offline reinforcement learning
batch reinforcement learning
deep reinforcement learning
combinatorial action spaces
structured action spaces
discrete action spaces
representation learning
Primary area: reinforcement learning

Abstract: Reinforcement learning in combinatorial action spaces requires searching over exponentially many joint actions to simultaneously select multiple sub-actions that form coherent combinations. Existing approaches either simplify policy learning by assuming independence across sub-actions, which often yields incoherent or invalid actions when coordination is required, or attempt to learn action structure and control jointly, which is slow and unstable. We introduce Structured Policy Initialization (SPIN), a two-stage framework that first pre-trains an Action Structure Model (ASM) to capture the manifold of valid actions, then freezes this representation and trains lightweight policy heads for control. On challenging DM Control benchmarks, SPIN improves average return by up to $39\%$ over the state of the art while reducing time to convergence by up to $12.8\times$.

Published as a conference paper at ICLR 2026
IMPROVING AND ACCELERATING OFFLINE RL IN LARGE DISCRETE ACTION SPACES WITH STRUCTURED POLICY INITIALIZATION
Matthew Landers'; Taylor W. Killian", Thomas Hartvigsen', Afsanch Doryab' University of Virginia, MBBUAT
ABSTRACT
Reinforcement learning in discrete combinatorial action spaces requires searching over exponentially many, joint actions to simultaneously select multiple sub-actions that form coherent combinations. Existing approaches either simplify policy learn- ing by assuming independence across sub-actions, which often yields incoherent or invalid actions, or attempt to learn action structure and control jointly, which is slow and unstable. We introduce Structured Policy Initialization (SPIN), two-stage framework that first pre-trains an Action Structure Model (ASM) to capture the manifold of valid actions, then freezes this representation and trains lightweight policy heads for control. On challenging discrete DM Control benchmarks, SPIN improves average return by up 39% over the state of the art while reducing time to convergence by up to 12.8
INTRODUCTION
Many real-world problems require decision-making in tigh-dimensional discrete action spaces, including applications in healthcare (Liu et al.) (2200) robotic assembly (Driess et al. 2020).. rec- ommender systems (Zhao et al. 2018). and ride-sharing (Lin et al. 2018). In such domains, online exploration can be costly or unsafe, making offline reinforcement learning (RL) (Lange et al. 2012
Levine et al. 2020) an appealing framework. Standard offline RL methods (Fujimoto et al. 2019
Agarwal et al. 2020 Fu et al. 2020 Kumar et al., 2020; Kostrikov et al. 2021)) however, are
not designed for large discrete action spaces, as they require either maximizing Q-function or parameterizing policy over the full discrete action set operations that become intractable the space scales exponentially with I[ my, where the number of sub-action dimensions and ma is the number of choices per dimension.
of Asymptotic Score
Learning in these complex settings requires solv-
100
ing two related problems: (i) searching over an exponential number of joint actions, and (ii) en-
suring that the chosen sub-actions form coherent combinations. Methods designed for these com- binatorial spaces have traditionally simplified policy learning by imposing strong structural priors such as assuming conditional indepen- dence between sub-actions (Tang et al. 2022) Beeson al. 2024). However, this sacrifices representational capacity, precluding the model from capturing the interactions required for ef- fective control. Other approaches attempt to learn the action representation and optimize policy simultaneously (Zhang et al. [2018] Lan-
50-
SPIN
Factored IQL
25
SAINT
AR-IQL
ASM Leamingg Complete
50
100
150
200
250
300
Time Minutes
In the humanoid d-ssdndd task. learning from medi um um-expert dataset, SPIN reaches the target performance (dashed horizontal line) quickly af- ter pre-training (dashed vertical line), while haselines require over 200 minutes of wall-clock training.
Figure 1:
ders et al.) 2024 2025). ). but this conflation of objectives often makes learning slow and unstable.
"wp4pk @ virginia.edu
Code available at httphhh::/gtthhb...oooo.. http:////g..ttccccoo/////aaaaaee

---

Published as a conference paper at ICLR 2026
We introduce Structured Policy INitialization (SPIN), two-stage framework that decouples repre- sentation learning from control. In the first stage, an Action Structure Model (ASM) is trained with self-supervision to learn representation function that, conditioned space over sub-actions in which structurally coherent joint actions concentrate on low-dimensional manifold. This action space representation is then frozen during the second stage, where the control problem reduces to learning lightweight policy heads over the action manifold for the downstream RL task. By learning structure first and policy second, SPIN allows the agent exploit the underlying action geometry instead of searching the raw combinatorial space. This leads to faster training and improved policy performance (Figure [T]. Across diverse benchmarks varying in dataset size and quality, action dimensionality, and action cardinality, SPIN improves average return by up to 39 % over the state the art and reduces training time to state-of-the-arr performance by up to 12.8x.
the state induces a feature
Our contributions are as follows:
We reframe offline RL in discrete structured action spaces as representation problem, separating learning action structure from control.
We propose SPIN, a two-stage framework that pre-trains and freezes an action-space repre- sentation to accelerate and improve policy learning.
We show that SPIN achieves state-of-the-art performance on challenging benchmarks, outperforming existing methods while being significantly faster.
We analyze the learned representations to demonstrate that capturing action structure is critical for effective policy learning in discrete combinatorial action spaces.
RELATED WORK
RL in Large Discrete Action Spaces. Several RL methods have been developed for combinatorial
action spaces in domains such routing Nazari et al. (2018); ; Delarue et al. (2020) and resource
allocation Chen et al. (2024). but these approaches typically rely on task-specific knowledge. General-
purpose methods have also been introduced Dulac-Arool et al. 2015 Tavakoli et al. 2018 Farquhar
et al. 2020 Van de Wiele et al. 2020; Zhao et al. 2023). but they are typically designed for online learning and are not easily adapted to the constraints of offline datasets. In offline RL, methods often factorize the policy or Q-function (Tang et al. 2022: Beeson et al. 2024). This factorization, however, enforces conditional independence across sub-actions, which restricts representational capacity and fails when sub-actions are strongly dependent. Other methods capture dependencies explicitly - BraVE (Landers et al.) (2224) models cross- cros-ddiensioonl interactions but scales poorly with action size, while autoregressive policies (Zhang et al. 2018) impose a fixed ordering that breaks permutation invariance. More recently, "AINT'(Lander
al. (025) introduced a Transformer- based policy to capture sub-action dependencies through self-attention, but learns action structure and control jointly, leading to slow and unstable training. A related line of work learns representations for large, but flat, action spaces. Most relevant is MERLION (Gu et al., 2022), which learns pseudometric- based action representation for offline RL. However, MERLION'S policy execution requires nearest-neighhor search over the full, enumerated action set every timestep, making it computationally infeasible for the combinatorial settings we consider. Furthermore, its architecture treats actions atomic entities and does not model their underlying compositional structure. SPIN, by contrast, is designed for this combinatorial regime, with a structured policy that generates joint actions dimension-win rather than enumerating the full combinatorial set.
Self-Supervised Pre-Training in RL. Self-supervised pre-training in RL has taken several forms
including as auxiliary objectives for representation shaping (Jaderbrr et al. (016)) Shelhhmer et al.
2016), contrastive and predictive encoders (Laskin et al., 2020; [Schwarzer et al. 2021
Stooke
et al. 2021 Liu & Abbeel, 2021bla), and world-modeling (Ha & Schmidhuber 2018). Other
work explores masked decision or trajectory modeling (Cai et | Liu et al. 2022; Wu et al.
2023 Sun et al.
2023). Large-scale behavioral pre-training has produced generalist policies and
vision language action models (Brohan et al. 2022 Zitkovich et al., 2023 O Neill et al.) 2024 Kim
et al. 2024 Team et al. 2024 Tirizzoni et al. 2025)_ with methods for rapid post-pre-tn pre-trainnnnnn adaptation
(Sikehi et al. (2025). These approaches are largely state/trajectory-cente and often presume online
interaction multi-task fine-tuning. SPIN, by contrast, pre-trains an Action Structure Model that
2

---

Published as a conference paper at ICLR 2026
captures action composition, providing structured initialization for policy learning in combinatorial action spaces without any online interaction.
PRELIMINARIES
An RL problem is formalized Markov Decision Process (MDP) M (,,A,, P, 7, 7,,), where S the state space, A the action space, p(s' a) the transition dynamics, r(s, a) the reward function, E [0. 1] the discount factor, and js the initial state distribution. A policy ι : S - (PAA) maps states distributions actions. The optimal policy maximizes the expected discounted return:
Σ 7rrs,, as) | ~ p.at x.. 8641 ~ p(: se. st.
arg HIXX E
Combinatorial action space. this work, assume actions compositional A A
The standard MDP formulation does not define structure in A. In
AN, with each component Aa a discrete set of size m.g An action a = (a1,.... aN) therefore lies in a space of exponential size |A| I ma, making naive maximization and policy parameterization infeasible.
Offline reinforvement learning. Offline reinforcement learning assumes access to fixed dataset B {(s,, de, Tı, $t+1) {= generated by behavior policy лT. The learning agent has addi- tional access the environment and must train entirely from this static dataset. In offline RL, the maximization in standard temporal difference learning:
L(0) Baa.aa.a"" (r max Q(d',(((d))) -(:.:)00)))
induces overestimation errors, since actions outside the support of B yield extrapolated val- ues Q(s', a'; 0-) without evidence. If such estimates are selected by the max operator, they propagated through Bellman updates. Unlike online RL, offline RL cannot correct these errors by interacting with the environment. Reliable offline learning therefore requires constraining policies to the support of I.
STRUCTURED POLICY INITIALIZATION (SPIN)
Structured Policy INitialization (SPIN) is two-stage framework for offline RL in structured action spaces that explicitly decouples representation learning from control. In the first stage, an Action Structure Model (ASM) is trained with self-supervision to learn a representation function that, conditioned the state 8, induces feature space over sub-actions in which structurally coherent joint actions concentrate on -iwddimnsiooal manifold. In the second stage, this representation is frozen, and policy learning is reduced training lightweight heads on the induced action manifold for the downstream RL task.
4.1 ACTION STRUCTURE MODELING (ASM)
SPIN's first stage trains Action Structure Model (ASM) that captures the structure of plausible actions, conditioned on the environment state. [ASM(s, o; v) denote Transformer encoder with parameters tt. The encoder operates on an input sequence X (MM+N)xd that concatenates M learned state embeddings (X........ Xssr with sub-action embeddings (Xa,...., Xax):
X (x.) X Na Xay) E (M++N)xd
We omit positional encodings across sub-actions to preserve petiuuaatioo-qquvvaiiininn over 01...., (Lee et al. 2019 Landers et al. 2025).
The ASM is trained with masked conditional modeling objective, analogous to masked language modeling in BERT (Devlin et al.) 2019). This objective enables the ASM to capture the manifold of valid actions directly from data, without requiring reward supervision. Because the ASM's role is to learn the manifold before policy optimization begins, this pre-training stage is fundamentally an offline procedure that requires static, pre-existing dataset. For each (s, a). we sample subset M N))..,N) of sub-action indices perturb. For every i M, a; is (i) replaced by mask

---

Published as a conference paper at ICLR 2026
Algorithm 1 ASM Pre-Training
Initialize parameters vi, heads {
for ( 1.... TasM do
Sample (8.a)~~ -0)0~~2 with = (a1s-..,ax)
Sample mask M (l,..., J; form by masking/replacing for M Compute attention over sub-actions: H fASM(s, " :t)
CASM DiCm Cross (f.(,) 0s)
Update 4. [ to minimize CASM
end for return //
token, (ii) replaced by an element drawn uniformly at random from the full sub-action space A,, or (iii) left unchanged, following an 80/10/10 ratio. The perturbed tuple maakk is then encoded by JASM- producing per-slot embeddings ha, at each slot Finally, each ha; mapped by slot-specific head ft : IR + IlA. to logits over A,, and cross-entropy loss computed only on the masked sub-actions:
AASM ((,a)~D BM Σ (J(hhaa): a:)
iEM
The ASM pre-training procedure is summarized in Algorithm | We empirically validate this objective in Appendix C where we show outperforms strong generative and discriminative alternatives.
4.2 POLICY LEARNING WITH A FROZEN REPRESENTATION
In the second stage, SPIN performs policy learning on the frozen representation provided by the ASM. The policy network To updates only lightweight components such as the query vectors and output heads, while the ASM remains fixed. This separation preserves the learned action structure and keeps policy optimization tractable.
Policy Architecture and Training SPIN implements the policy me using the SAINT architecture (Landers et al. 2025), which models dependencies among sub-actions while preserving tractability. Specifically, M state embeddings and N learnable action queries are passed through the frozen, cermutation-equivariian
Transformer from the ASM stage, producing contextualized embeddings 21,..., NN. Each embedding z; encodes the state, the corresponding action query, and its relations other sub-actions through shared attention. These embeddings are then passed to sub-action-specifiic MLP heads fi, which output logits & over the corresponding sub-actions. The logits parameterize categorical distributions (- 8, 2,) = softmax from which is sampled. The joint policy factorizes these contextualized distributions:
N
(a
II mo 8, 2,)
As in prior factored approaches (Tang al. 2022 Beeson al. 2024). this factorization preserves tractability rather than optimizing over an exponentially large joint action space, the policy produces N categorical distributions. Unlike purely factored methods that assume conditional independence across sub-actions. SPIN retains cross-dimensional dependencies through shared self- attention, learned during ASM pre- pre-training, to produce contextualized embeddings z; The full policy learning procedure is summarized in Algorithm 2]
Compatibility with Offline RL Methods. This design makes SPIN broadly compatible with actor-critic algorithms for which the actor update is expressed as weighted log-likelihood maxi- mization over dataset actions:
max E( (s. a) log me (a {s]
where Wip (s, a) 2 encodes algorithm speciffc weights (e.g., advantages or value estimates). This class includes methods such as IQL (Kostrikov et al. 1022 and AWAC (Nair et al.) 2020). SPIN also supports selection-basee updates, as in BCQ (Fujimoto et al. 2019), where the actor is trained on dataset-supported candidate actions.
4

---

Published as a conference paper at ICLR 2026
Algorithm 2 Policy Learning with Frozen Representation
Initialize policy params 0 (queries Q, heads), aux params 4 (c.g., critic), optimizers
for t 1...1...... .u. do
Sample o, r,
Get contextualized embeddings: B1:N f- (s, s, Q; ('tazen)
D Policy forward pass
Compute policy loss: La (w a) > log (a, 2.)
Compute auxiliary loss Le with chosen offline RL objective
Update both and & with their respective gradients
Only queries/heads train
end for return 0
Objectives requiring global operations over the joint action space are intractable unless Qs or ma factorized across action dimensions, which enforces conditional independence and discards cross-dimensional structure. This conflicts with SPIN's objective of modeling action coherence, so value-regularizzttion methods such as CQL (Kumar et al. 2020) fall outside this compatibility class, mirroring the boundary defined by SAINT.
5 EXPERIMENTAL EVALUATION
We evaluate SPIN discretized variant of the DeepMind Control Suite (Tassa et al. 2018), introduced by Beeson et al. (024). The combinatorial action spaces grow exponentially with the number of joints and the sub-action cardinalities. Action spaces vary along two axes: (i) the action dimensions, ranging from six in cheetah to 38 in trot and (ii) the sub-action cardinalities, ranging from three thirty bins per dimension. Together, these variations yield joint action spaces spanning several orders of magnitude, from hundreds to 30 x 1.35 x 10 possible actions.
Datasets are constructed at four quality levels following standard offline RL protocols. The medium sets are generated by partially trained policies, and the expert sets by fully trained policies. The medi um-expert sets combine transitions from both, while the random-med: wew-nnmmmmdu--e-eeeeeeeeeee um-expert sets additionally include trajectories with random actions, yielding highly heterogeneous distributions. Dataset sizes range from 2 10y to 2 x 10 transitions.
We evaluate SPIN against three baselines representing the primary approaches to offline learning in structured action spaces. SAINT (Landers et al., 2025) is Transformer-base policy architecture that models sub-action interactions through self-attention, jointly learning both the action structure and control policy under a single RL objective. our evaluation, SAINT serves as the most direct joint- learning comparison to SPIN. We instantiate both methods within the same policy class to ensure that differences in performance reflect the learning paradigm rather than architectural variations. Under this controlled setup. SAINT differs from SPIN precisely in the dimension our work investigates: the separation of representation learning from control. A factored policy (Tang et al., [202]] [eessn al. 2024 assumes conditional independence across sub-actions, learning separate per-dimension distributions without modeling interactions. An autoregressive policy (Zhang et al. 2018) models
dependeneies sequentially by factorizing the joint distribution as a chain, conditioning each sub- action on its predecessors. This autoregressive decomposition depends on an arbitrary ordering of dimensions and not permutation-equiiariaat.
To isolate the effect of architectural choices, all methods are trained with the IQL (Kostrikov et al. 202T) objective. To assess robustness, we also report results with alternative objectives, including AWAC (Nair et al. 2020) and BCQ (Fujimoto et al., 2019), in Appendix [ To validate SPIN's generalizability beyond locomotion, we evaluate its performance on Maze (Beeson et al.) 2024), with results provided in Appendix E To demonstrate that SPIN's effectiveness is due its
pre-training objective rather than from pre-training alone, we compare its performance to that of trajectory-centric pre-training approach in Appendix [F] Across all of these settings, SPIN consistently outperforms the baselines both performance and efficiency.
All experiments using Python 3.9 with PyTorch 2.6 single NVIDIA A40 GPU. Reported results are averaged over five random seeds, with + values indicating one standard deviation across seeds.
5

---

Published as a conference paper at ICLR 2026
Task
F-IQL
AR-IQL
SAINT
SPIN
Medium
cheetah
293.0 6.9
2.4
6.1
293.6 4.1
finger
385.4 6.6
383.5
3.7
392.7
humanoid
335.4 + 6.0
327.3
$82.3=722
334.8 7.6
quadruped
353.4 + 82.4
3134 814
3
359.7 + 78.7
Average Return
341.8
334.7
343.0
345.2
Time to Target
48.2
114.3
174.6
45.5
Medium-Expent cheetah
612.9 + 50.2
609.9 37.7
37.
651.1 33.1
finger
844.5 =
8
847.6
humanoid
603.0 + 49.5
567.6 + 50.4
621.5
652.5 31.0
quadruped
838.2 + 45.9
833.4 + 46.4
836.5 + 35.9
854.1 + 37.7
Average Return
724.7
717.2
733.3
753.2
Time to Target
257.3
285.8
308.4
62.0
cheetah
289.6 + 17.0
276.5 13.2
3 33.2
332.4 43.1
finger
693.4 + 17.5
762.7 20.1
747.2 17.1
773.2 + 14.9
humanoid
230.1 =
192.8 24.1
29.4 237.3
+ 28.6
quadruped
340.0 = 52.6
350.2 = 84.5
468.8 55.5
561.0 + 70.6
Average Return
388.3
395.6
438.9
499.2
Time to Target
85.1
95.8
100.2
38 4
Expert
cheetah
665.9 + 24.2
665.7 + 20.1
672.7 1: 21.5
finger
874.0 + 4.6
873.3 1 6.7
$70.219.2
868.5 1 6.4
humanoid
729.0 21.2
717.2 + 25.0
726.1 25.3
734.7 + 19.5
quadruped
843.5 20.2
824.9 + 35.8
831.6 43.9
839.0 + 33.9
Average Return
778.1
770.3
773.1
778.7
Time to Target
167.5
288.7
261.8
77.4
Overall Average Return
558.2
551.5
572.1
594.1
Overall Time Target
558.1
784.6
845.0
223.3
Table 1: Asymptotic performance and training efficiency DM Control tasks. Time to Target denotes the wall-elock minutes required to reach 95% of F-IQL's asymptotic performance, with SPIN times including ASM pre-training. SPIN attains the best overall returns and the highest computational efficiency.
5.1 ASYMPTOTIC PERFORMANCE AND TRAINING EFFICIENCY
Table T reports final performance and training efficiency across environments and dataset qualities (full learning curves are provided in Appendix A). SPIN achieves consistently higher returns than all baselines and reaches the target performance in less wall-clock time than all baselines.
SPIN achieves the highest overall average return of 594.1, exceeding the next-best baseline, SAINT, 572.1. Improvements are systematic across the benchmark suite rather than concentrated in individual environments. The advantage most pronounced in the heterogeneous medi um expert and random-nediim-- uim-d-iiuureeeeee datasets, which represent the most realistic and challenging benchmark settings. On the random-mes wwt-d-mmmeiii----eeeeee--------e un-expert datasets, SPIN achieves an average return 499.2. 2, an improvement of more than 13% over the next-best method, SAINT (438.9).
We also measure the wall-clock time, reported as the number of minutes, required for each method to reach 95% of F-IQL's asymptotic performance. F-IQL is widely adopted state-of- ooatte--f-hh-att the-art art baseline in structured action spaces (Tang et al. 2022 Beeson et al., 2024 Landers et al. (024), offering both tractability and stable convergence across enviromments. Using F-IQL as the target enables fair comparison across methods that converge to different return levels, avoiding misleading advantages from terminating early at suboptimal performance. We adopt the 95% threshold rather than 100% because some methods never reach F-IQL's asymptotic performance. Handling these cases directly either by excluding runs by reporting full runtimes would bias the results, whereas the 95% %

---

Published as a conference paper at ICLR 2026
Bins
F-IQL
AR-IQL
SAINT
SPIN
Bins
472.3 1 43.5
526.5 57.8
635. 4: 39.6
647.0 19.6 629.5 52.5
Bins
483.8 + 33.0
457.4 53.2
529.1 58.2
30 Bins
485.0 = 54.0
557.4 66.0
562.5 88.7
703. 99 25.6
Average Return
480.4
513.8
575.6
660.1
Time to Target
545.8
692.6
291.6
237.0
Table 2:
Performance and efficiency on the dog-trot task action cardinality increases. SPIN sustains strong returns and training efficiency as the action space grows, while basclines stagnate deteriorate.
criterion offers a consistent and comparable metric. Full per-environment time-to-targe results are reported in Appendix total, SPIN reaches the target performance 223.3 minutes, approximately 2.5x faster than F-IQL. itself and 3.8 faster than SAINT. The acceleration is especially pronounced the medí um-expert datasets, where SPIN requires only 62 minutes training time compared to more than 250 minutes for all other methods. All runtimes for SPIN include the full duration of the ASM pre-training stage.
These findings demonstrate that explicitly modeling action structure in a dedicated pre-training phase allows the representation to capture the manifold of coherent actions. Freezing this representation during policy learning preserves that structure, enabling lightweight heads to adapt efficiently to the downstream task. Compared to Factored and Autoregressive approaches, which either discard or impose rigid structure cross diosedimnniional dependencies, SPIN retains flexibility without sacrificing tractability. Unlike SAINT, which attempts to learn action structure and control jointly, SPIN's decoupled design achieves both higher asymptotic performance and faster convergence.
5.2 ROBUSTNESS ACTION CARDINALITY
The results in Section 5.1 report evaluations with fixed action cardinality of three bins per dimen- sion. To test robustness under more severe combinatorial growth, we increased the granularity of discretization in the dog-t tro environment, which has 38 sub-action dimensions. Varying the cardinality from to 30 bins produces action spaces ranging from 3 1.35 x 1028 to more than 3038 1.35 x 1036 possible actions. Experiments which provides a realistic and challenging setting.
conducted the medi um exprrt dataset,
Results are summarized in Table 2 SPIN achieves the highest average return at every cardinality, and the gap relative to baselines increases with the size of the action space. At three bins, SPIN slightly outperforms the strongest baseline, SAINT. At thirty bins, SPIN reaches average return of 703.9 9 compared to 562.5 for SAINT, an improvement of more than 25%. AR-IQL shows unstable performance, dropping from 526.5 three bins 457.4 from increased granularity, plateauing around 480.
bins, while F-IQL shows no benefit
Training efficiency follows the same trend. SPIN consistently requires less wall-clock time to reach the target performance, even in the largest action spaces (full runtime results are provided in Appendix B These results demonstrate that separating structure learning from control is increasingly beneficial as combinatorial complexity grows, as the agent can act over learned now-dimensional manifold while end-to-end methods remain tied to the scale of the raw joint space.
MECHANISMS UNDERLYING SPIN'S EFFECTIVENESS
The experiments in Section 5 show that SPIN outperforms existing methods both in learning speed and final performance, We now examine the mechanisms underlying these gains.
6.1 EFFECT OF REPRESENTATION QUALITY ON POLICY PERFORMANCE
To assess the contribution ASM pre-training, we trained the ASM representation on the medi um-expert datasets for 10-100 epochs. Each representation function then frozen and used to initialize a new policy, which was subsequently trained convergence on the control task. Figure shows that downstream return generally improves with ASM pre-training, with the
7

---

Published as a conference paper at ICLR 2026
Cheetah
Finger
Humanoid
Quadruped
1pa3
7pph
1pms
1pock
Figure 2: Final policy return as function of ASM pre-training duration. Policies are initialized from frozen ASMs trained for different numbers of epochs. While pre-training is critical, just 20 epochs yield a representation that enables policies to surpass the fully trained F-IQL baseline in all tasks (dashed horizontal line).
steepest gains in the first 20 epochs. After 20 epochs, policies surpass the fully converged F-IQL, refer- ence all tasks. Policies initialized from untrained ASM (Epoch 0) perform poorly. These results indicate that final policy performance is largely determined by the quality of the pre-trained action representation; once coherent representation is learned, control optimization becomes substantially easier.
6.2 QUANTIFYING REPRESENTATION QUALITY
The large gap between randomly initialized (epoch 0) and pre-trained agents in Figure may be due to pre-training providing only a convenient initialization, without encoding structure, or due to pre-training learning a representation that enables downstream performance. We evaluate this directly by testing whether the ASM representation captures joint action dependencies using linear probe, standard diagnostic for self-supervised representations (Chen et al.) 2020 He et al. 2020).
In this experiment, the ASM representation is frozen either pre-trained for 100 epochs randomly initialized - and a lightweight linear classifier is trained on its embeddings to predict dataset actions from the state. New action queries and linear heads are learned for this probe. The analysis is conducted in the dog-trot environment, which has 38 sub-action dimensions discretized into 30 bins, yielding the largest and most challenging combinatorial action space in the DM Control suite.
The random ASM (0 epochs) achieves 76.6% per-slot accuracy, but its exact-match accuracy on the full 38-dimensional action tuple is only O. 1.10% well above the uniform-chance baseline of 30-38 R 7 x 10-57, but far below the level required for coherent control. These features capture individual action dimensions but fail to encode cross-joint dependencies. The pre-trained ASM, by contrast, attains 90.0% per-slot accuracy and 4.52% exact-match accuracy, a 45x improvement over the random ASM. Crucially, the observed tuple accuracy more than doubles the value expected under independence (0.90 x 1.83%). showing that pre-training produces representation with substantial cross-joint coordination.
6.3 ISOLATING THE CONTRIBUTION THE LEARNED REPRESENTATION
The analyses in Sections 6.1 and 6.2 indicate that final policy performance is largely determined by the quality the pre-trained action representation. This suggests a sufficiently expressive representation should reduce the architectural demands on the downstream policy, allowing strong performance to be achieved even with simple policy head. To test this hypothesis directly, we introduce SPIN-DistilL, variant that replaces the Trunsformer-basee policy in Stage (Section 4.2 with lightweight, attention frreeinn-feeeeeee MLP. Specifically, the knowledge from the frozen ASM fASM 5, Q: e'trozen is distilled into student network, (s, Q) using mean-squared error objective reproduce the contextualized embeddings of the teacher:
Drisill = EEED [1a0(s, l,s5,,00 Q) - SASM[(S, (8, Q; virrnen))
)|
After training, the student network is frozen and functions as a lightweight, attention-fiee feature extractor for the downstream policy.
the results of this experiment.
SPIN-Distill comes within a small margin of the full SPIN model's asymptotic performance and substantially outperforms all other baselines, while being nearly 8x faster than SAINT. These results provide strong evidence that SPIN's performance gains are attributable to the quality of the pre-trained representation itself, and not the policy network's 's specific architecture.
8

---

Published as a conference paper at ICLR 2026
Task
F-IQL
AR-IQL
SAINT
SPIN
SPIN-Distill
cheetah
612.9 + 50.2
609.9 37.7
627.4-37.J
631.1113..1
645.3 1 40.8
finger
844.5 = 11.1
857.8
847.6
8.5.2
852. 9.3
humanoid
603.0 = 49.5
567.6 50.4
621.5 53.5
31.0
650.9 + 28.8
quadruped
838.2 = 45.9
833.4 + 46.4
836.5 + 35.9
854.1 37.7
846.8 + 27.7
Average Return
724.6
717.2
733.3
753.2
748.8
Time to Target
257.3
285.8
308.4 4
62.0
39.2
Table 3: Performance on medi um-expert datasets. SPIN-Distill achieves asymptotic performance similar SPIN, indicating that a strong learned representation can compensate for a simpler policy architecture.
6.4 EMERGENT RAPID ADAPTATION
Having established the importance of pre-training and representation quality, we next examine learning dynamics. Table reports the percentage of F-IQL's asymptotic performance achieved after 10,000 gradient steps, corresponding to only 1% of the total training budget. Across nearly all environments, SPIN learns policies that reach at least 90% of target performance, whereas baselines improve much more gradually. The effect is most pronounced on heterogeneous datasets. In the humanoic d task with the medi um expert dataset, SPIN reaches 93.4% of target performance, while the next-best method, SAINT, achieves only 9.3%. On the random nandom-medi medi um expert datasets. SPIN exceeds 100% of F-IQL's asymptotic performance in both cheetah and humanoid during this period.
This rapid learning also clarifies SPIN's wall-clock efficiency (Table D. The downstream RL stage is computationally dominated by the actor-critic loop, which requires repeated evaluations of the actor, critic, and target networks as well as Bellman backups at every gradient step. The ASM pre-training stage, by contrast, is stable, single-pass supervised objective applied over masked sub-actions. Its relative cost is therefore minimal: on the medi um expert datasets, pre-training constitutes only 5.6% of total wall-clock time on cheetah, 1.4% on f finger, and 1.6% both humanoic d and quadruped.
Together, these results suggest that the ASM provides strong structural prior that substantially simplifies downstream learning. End-to-end baselines must jointly discover both action structure and control, leading to slow initial progress, whereas SPIN begins policy learning with coherent representation, enabling efficient early adaptation and reduced overall training time.
7 DISCUSSION AND CONCLUSION
Reinforcement learning in discrete combinatorial action spaces requires searching over an exponential number of composite actions while ensuring that the chosen sub-actions form coherent sets. Some
methods simplify policy learning by ignoring action structure (Tang et al.) 2022 Beeson et al. 2024), the cost of discarding critical sub-action dependencies. Other approaches attempt pt simultaneously 2025), but are often prohibitively slow and unstable. SPIN, by contrast, separates representation learning from policy learning using a two-stage procedure. In the first stage, an Action Structure Model (ASM) learns a representation function that, conditioned on the state s, induces feature space over sub-actions in which structurally coherent joint actions lie on oow-dimensional manifold. This representation then frozen and reused in the second stage, where control reduces to training lightweight policy heads on top of the pre-trained ASM.
capture structure and solve control (Zhang et al.) 22188) Landers et al., 2024
Across benchmarks varying in dataset size and quality, action dimensionality, and action cardinality, SPIN improves average return by up to 39% over the state of the art and reduces the time to reach strong baseline performance by up to 12.8x. Gains are most pronounced the challenging. realistic medi um-expert and r random-med um-expert datasets.
targeted analysis elucidates SPIN's s effectiveness. Final performance rises with the quality of the learned representation, confirming that control is bottlenecked by structure discovery. Once this structure is available, policies learn rapidly, reaching most of their eventual return within small fraction of training. Linear probes further show that the learned representation is 45x effective

---

Published as a conference paper at ICLR 2026
Task
F-IQL
AR-IQL
SAINT
SPIN
Medium
cheetah
94.0%
91.0%
94.6%
94.5%
finger
90.2%
89.0%
89.8%
91.2%
humanoid
50.9%
67.3%
84.9%
94.2%
quadruped
90.3%
82.7%
92.7%
94.0%
Average
81.4%
82.5%
90.5%
93.5%
Medium- Expert
cheetah
30.3%
40.9%
46.1%
90.6%
finger
0.37%
0.70%
0.39%
68.5%
humanoid
2.2%
3.1%
9.3%
93.4%
quadruped
34.1%
38.1%
39.0%
78.6%
Average
16.7%
20.7%
23.7%
82.8%
Kandom-Mediune
Expert
cheetah
87.7%
93.3%
94.0%
103.6% 6%
finger
31.1%
34.6%
33.6%
49.0%
humanoid
23.3%
30.4%
50.5%
110.7%
quadruped
43.3%
48.2%
45.0%
80.7%
Average
46.4%
51.6%
55.8%
86.0%
Expert
cheetah
26.5%
35.8%
49.1%
94.3%
finger
31.6%
33.1%
8.4%
94.9%
humanoid
2.8%
2.3%
14.0%
95.1%
quadruped
44.6%
45.7%
53.0%
79.2%
Average
26.4%
29.2%
31.1%
90.9%
Overall Average
42.7%
46.0%
50.3%
88.3%
Table 4: Percentage of final F-IQI. performance achieved after 10,000 gradient steps (1 % of training budget). SPIN quickly learns high-performing policy, while haselines improve more gradually.
at producing fully coordinated actions than a random baseline, providing direct, quantitative explanation for the downstream agent's success.
While SPIN demonstrates strong performance, several opportunities for future work remain. Ex- tending SPIN reuuaaaagulii methods such as CQL is promising direction. One natural next step is to develop hybrid objectives that combine SPIN's representation-firs design with mild conservative regularization actions or applied at the sub-action level, thereby avoiding intractable global operations over the full combinatorial space. Adapting SPIN action spaces that exhibit structural assumptions other than permutation equivariance, such as ordered sequence- based sub-actions is another direction for future work. Finally, as with all offline methods, SPIN's generalization ultimately depends on dataset coverage, and improving robustness under sparse or biased data remains an important open challenge.
for example, penalties restricted ASM-proposed candidate joint
SPIN introduces representation-fine ffrstt view of control in structured action spaces. By first learning the manifold of plausible actions and then reusing representation function for downstream decision- making, reduces complex combinatorial problem to tractable policy learning task. This decoupling offers principled framework for reinforcement learning high-dimensional, structured domains.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International conference on machine learning, pp. 104-114. PMLR, 2020.
Alex Beeson, David Ireland, and Giovanni Montana. An investigation of offline reinforcement learning in factorisable action spaces. arXiv preprint arXiv:2411.11088, 2024.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:22/2.068// 2022.
Yuanying Cai, Chuheng Zhang, Wei Shen, Xuyun Zhang, Wenjie Ruan, and Longbo Huang. Reprem: Representation pre-training with masked model for reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 6879-6887 2023.
Changyu Chen, Ramesha Karunasena, Thanh Nguyen, Arunesh Sinha, and Pradeep Varakantham. Generative modelling stochastic actions with arbitrary constraints in reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PmLR, 2020.
Arthur Delarue, Ross Anderson, and Christian Tjandraatmadja. Reinforcement learning with combi- natorial actions: An application to vehicle routing. Advances in Neural Information Processing Systems, 33:609-620, 2020.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171 1181-488,, 2019.
Danny Driess, Jung-Su Ha, and Mare Toussaint. Deep visual reasoning: Learning predict action sequences for task and motion planning from an initial scene image. arXiv preprint arXiv:2006.05330 05398, 2020.
Gabriel Dulac-Arnold, Richard Evans, Hado van Hasselt, Peter Sunehag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep reinforcement learning in large discrete action spaces. arXiv preprint 1rX:::1112.07779 2015.
Gregory Farquhar, Laura Gustafson, Zeming Lin, Shimon Whiteson, Nicolas Usunier, and Gabriel Synnaeve. Growing action spaces. In International Conference on Machine Learning, pp. 3040- 3051. PMLR, 2020.
Yuwei Fu, Wu Di, and Benoit Boulet. Batch reinforcement learning in the real world: A survey, In Offline RI. Workshop, NeurolPS, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference machine learning, pp. 2052-2065 PMLR, 2019.
Pengjie Gu, Mengchen Zhao, Chen Chen, Dong Li, Jianye Hao, and Bo An. Learning pseudometric- based action representations for offline reinforcement learning. 2022.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 180::8880.1112.. 2(3). 2018.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint 161 1.03997, 2016.
11

---

Published as a conference paper at ICLR 2026
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, et al. Openvla: An open-source model. arXiv preprint arXiv:2406.089 00200000992.... 2024.
vision- Aisionlanguuage-e
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Iya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169, 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:1179-111 2020.
Matthew Landers, Taylor W Killian, Hugo Barnes, Thomas Hartvigsen, and Afsaneh Doryab. Brave: Offline reinforcement learning for discrete combinatorial action spaces. arXiv preprint arXiv:2410.21154, 2024.
Matthew Landers, Taylor W Killian, Thomas Hartvigsen, and Afsanch Doryab. Saint: Attention-based modeling of sub-action dependencies in multi-action policies. arXiv preprint arXiv:5505.12109, 2025.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-ar,, pp. 45-73. Springer, 2012.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650. PMLR, 2020.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for attention-based permutation-invarian neural networks. In International conference on machine learning. pp. 3744-3753. PMLR, 2019.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01644 2020.
Kaixiang Lin, Renyu Zhao, Zhe Xu, and Jiayu Zhou. Efficient large-scale fleet management via multi-agent deep reinforcement learning. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1774-1783, 2018.
Fangchen Liu, Hao Liu, Aditya Grover, and Pieter Abbeel. Masked autoencoding for scalable and generalizable decision making. Advances in Neural Information Processing Systems, 35: 16608-12618 2022.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. International Conference on Machine Learning, pp. 6736 6667-6747. PMLR, 2021a.
Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 34:18499-18473, 2021b.
Siqi Liu, Kay Choong See, Kee Yuan Ngiam, Leo Anthony Celi, Xingzhi Sun, and Mengling Feng. Reinforcement learning for clinical decision support in critical care: comprehensive review. Journal of medical Internet research. 22(7):e 1877, 2020.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09355 2020.
Mohammadreza Nazari, Afshin Oroojlooy, Lawrence Snyder, and Martin Takác. Reinforcement learning for solving the vehicle routing problem. Advances in neural information processing systems. 31, 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: / 1807.0877733 2018.
12

---

Published as a conference paper at ICLR 2026
Abby O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, et al. Open x-embodiment: Robotic learning datasets and rt-x models: Open x-embodiment collaboration O. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 6892-6903. IEEE, 2024.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations, 2021.
Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell. Loss is its own reward: 2016. ICLR 2017
Self-supervision for reinforcement learning. arXiv preprint Workshop.
Harshit Sikchi, Andrea Tirinzoni, Ahmed Touati, Yingchen Xu, Anssi Kanervisto, Scott Niekum, Amy Zhang, Alessandro Lazaric, and Matteo Pirotta. Fast adaptation with behavioral foundation models. arXiv preprint rrXivv2504.078%. 2025.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9870-9879, PMLR, 2021.
Yanchao Sun, Shuang Ma, Ratnesh Madaan, Rogerio Bonatti, Furong Huang, and Ashish Kapoor. Smart: Self-supervised multi-task pretraining with control transformers. arXiv preprint arXiv:2301.09816 2023.
Shengpu Tang, Maggie Makar, Michael Sjoding, Finale Doshi-Velez, and Jenna Wiens. Leveraging factored action spaces for efficient offline reinforcement learning in healthcare. Advances in Neural Information Processing Systems, 35:34272-342885 2022.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXhv:1801.00690, 2018
Arash Tavakoli, Fabio Pardo, and Petar Kormushev. Action branching architectures for deep rein- forcement learning. In Proceedings of the aaai conference on artificial intelligence, volume 32, 2018
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprint a40i:: 12213, 2024.
Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot whole-body humanoid control via behavioral foundation models. arXiv preprint arXiv:2504.1 1004,, 2025.
Tom Van de Wiele, David Warde-Farley, Andriy Mnih, and Volodymyr Mnih. Q-learning in enormous action spaces via amortized approximate maximization. arXiv preprint arXiv:2001.08/ 16, 2020.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. Masked trajectory models for prediction, representation, and control. International Conference on Machine Learning, pp. 37607 PMLR, 2023.
Yiming Zhang, Quan Ho Vuong, Kenny Song, Xiao Yue Gong, and Keith W Ross. Efficient entropy for policy gradient with multidimensional action space. arXiv preprint / 2018.
Xiangyu Zhao, Long Xia, Liang Zhang. Zhuoye Ding, Dawei Yin, and Jiliang Tang. Deep reinforce- ment learning for page-wise recommendations. In Proceedings of the 12th ACM conference on recommender systems, pp. 95-103, 2018
Xutong Zhao, Yangchen Pan, Chenjun Xiao, Sarath Chandar, and Janarthanan Rajendran. Condition- ally optimistic exploration for cooperative deep multi-agent reinforcement learning. In Uncertainty in Artificial Intelligence, pp. 2529- 2540- PMLR, 2023.
13

---

Published as a conference paper at ICLR 2026
Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, et al. Rt-2: Vision- Usiio-languagg-actiinn models transfer web knowledge to robotic control. In Conference Robot Learning, pp. 2165-2183 PMLR, 2023.
14

---

Published as a conference paper at ICLR 2026
A LEARNING CURVES
Figure 3 shows the full learning curves corresponding to the results in Table |
Each plot reports
mean episode return versus gradient steps for all methods across environments and dataset qualities.
SPIN
F-IQL
SAINT
AR-IQL
Cheetah ah
Finger
Humanoid
Quadruped
medium
XII
254
Osstralthpp bkgs
Coukca 3kg
Oualaa 3kg
Grubben Nip
medi um-expert
389
289
13
421
13
421
3kg
S0p
random-med www.mmedddddeeeeeeeeee lum-expert
Bereen
100
200
135
208
481
134
480
Skys
1kg
expert
Donus
200
Goddar 3kg
Gralaa 3kgp
Qradiar Stiee
Figure 3: Learning curves on all DM Control tasks and dataset qualities. Curves show mean return over 1M gradient steps, averaged across five seeds with shaded regions denoting 11 std.
The curves visually confirm the trends reported the main text SPIN matches or exceeds the best baseline policy in every environment. SPIN's advantage is most pronounced the medi deiummexpeet and randon-medd um-expert datasets, the most challenging and realistic settings.
15

---

Published as a conference paper at ICLR 2026
B WALL-CLOCK EFFICIENCY
This section reports the full wall-clock training times underlying the efficiency analyses in Sections 5.1 a.d For comparability across methods with different asymptotic returns, measure the time required to reach 95% of the Factored IQL (F-IQL) baseline's 's final performance. We adopt the 95% target rather than 100% because all methods reach F-IQL's asymptote. Alternttiee treatments such as excluding these runs or reporting their full runtime would bias the averages. The 95% criterion provides consistent point of comparison across methods. All reported SPIN runtimes include the complete ASM pre-training stage.
B. WALL-CLOCK TRAINING TIMES FOR SECTION 5.1
Task
F-IQL
AR-IQL
SAINT
SPIN
Medium
cheetah
7.0
13.0
12.3
2.4
finger
4.7
8.0
11.7
7.8
humanoid
25.1
59.3
43.1
23.5
quadruped
11.4
34.0
107.5
11.8
Medium Total
48.2
114.3
174.6
45.5
Medium-Expert
cheetah
80.0
84.1
103.3
14.9
finger
28.8
20.5
56.3
14.0
humanoid
128.2
143.2
112.2
11.2
quadruped
20.3
38.0
36.6
21.9
Medium-Exper Total
257.3
285.8
308.4
62.0
Random-Medium--Bxy
cheetah
11.4
13.6
17.9
3.0
finger
33.5
10.6
33.6
25.7
humanoid
23.7
55.6
33.1
5.6
quadruped
16.5
16.0
15.6
4.1
Random-Medium- Total
85.1
95.8
100.2
38.4
Expert
cheetah
27.6
46.3
53.3
19.6
finger
16.5
16.9
38.5
4.3
humanoid
110.0
199.6
147.3
45.6
quadruped
13.4
25.9
22.7
7.9
Expert Total
167.5
288.7
261.8
77.4
Total Runtime
558.1
784.6
845.0
223.3
Table Wall-clock training time (minutes) to reach 95% of F-IQL's asymptotic performance, with SPIN times including ASM pre-training. Italicized rows give totals per dataset quality; the bottom row reports the overall total. This table provides the detailed breakdown for the Time to Target results in Table
Table Table repoors the per-environment efficiency results summarized in Table T SPIN is consistently the fastest method, with the largest gains on the med i um com-expert and random-men med um-expert datasets.
16

---

Published as a conference paper at ICLR 2026
B.2 WALL-CLOCK TRAINING TIMES FOR SECTION 5.2
Bins
F-IQL
AR-IQL
SAINT
SPIN
3 Bins
388.3
319.3
121.7
81.2
10 Bins
102.1
215.2
87.9
77.8
30 Bins
55.4
158.1
82.0
78.0
Time Target
545.8
692.6
291.6
237.0
Table 6: Wall-clock training time (minutes) to target performance in the dog-t roo environment as action cardinality increases. This table provides the detailed breakdown for the Time to Target results in Table
Table reports detailed runtimes for the action cardinality experiments in the dog-t environ- ment, corresponding to Table 2 SPIN is the most efficient method, achieving the lowest total training time across all experiments. Overall, SPIN is the fastest 3 and 10 bins, while the simpler F-IQL model converges quickest at 30 bins. However, F-IQL's speed comes at the cost of much lower final performance, as shown in Table Crucially, SPIN's runtime remains stable even as the action space expands by many orders of magnitude, demonstrating superior scalability. These results confirm that SPIN offers the best overall balance of strong performance and computational efficiency.
17

---

Published as a conference paper at ICLR 2026
C COMPARISON TO ALTERNATIVE PRE-TRAINING OBJECTIVES
To isolate the contribution of the pre-training objective, we compare SPIN's masked modeling to that of strong baselines from the other dominant paradigms in self-supervised learning: generative objective (Variational Action Modeling) and a discriminative objective (Contrastive Action Modeling). All experiments are rum on the challenging and realistic medi um-expert datasets.
C.I METHODOLOGY
Masked Action Modeling (MAM). This is the reconstructive objective used by SPIN, as described in detail in Section 4.1
Variational Action Modeling (VAM). As generative alternative, adopt Variational Action Model (VAM) that learns a probabilistic latent representation of plausible actions, The model com- prises Transformer-based encoder gz(2|s,
ss a) and an MLP decoder po(als,2). The encoder reuses the ASM Transformer fASM (s, a; 45) to produce contextualized slot embeddings (hat:.... hav ). which mean-pos mean-pooled to form global representation h, t Σ ha. This representation is mapped to the parameters of diagonal Gaussian posterior, Po (h,) and log (h. from which latent code z E &R is sampled using the reparameterization trick (Kingma & Welling 2013). The decoder po reconstructs the action by mapping (concatenated with state embedding) logits for each sub-action dimension via set of parallel, slot-specific heads (ff Rd.(+d.) + RI4. Jim1 The model trained by maximizing the Evidence Lower Bound (ELBO):
N - [ll()))
LVAM = (s (xo(als.))
N
where p(z) = (0, II is standard normal prior and / balances reconstruction and regularization.
Contrastive Action Modeling (CAM). We further consider a discriminative objective that learns action structure through contrastive prediction. The Contrastive Action Model (CAM) uses the same Transformer encoder fASM(s, a; v)), trained with a Local-to-Global InfoNCE loss (Oord et al. 2018). The objective encourages each sub-action representation ("local"") to be predictable from the joint representation of the remaining sub-actions ("global"). For an unmasked action (s, a), slot embeddings (hay...., bay)
first computed. For each slot the local token is ha; and the global context is the mean of the others, hax.i = N-T Enz; bay. These are projected through separate heads Dokk (-) and 9ax(-) (-) to yield normalized embeddings qi and k, The model is trained to maximize agreement between matching (q,,k, pairs while minimizing agreement with negatives, defined as the context embeddings {kg Je from all other action tuples in the batch. The InfoNCE loss for slot is defined as:
B
exp(g
kg /T)
LCAM. =
log
B
exp(q.
K (3) /T)
b=1
=1
where 7 is temperature hyperparameter (Chen et al. 2020). The overall loss averages across all sub-action positions:
LCAM
CCAM.
18

---

Published as a conference paper at ICLR 2026
C.2 EMPIRICAL COMPARISON
Task
MAM
VAM
CAM
cheetah
651.1 + 33.1
607.5 + 54.2
607.3 + 46.6
finger
855.2 9.7
+1.5
851.3 = 15.9
humanoid
652.5
638.5 39.0
39.0
quadruped
854.1-37.1
845.5
838.5 = 44.7
Table 7: Mean + std policy performance medium-expert tasks under different ASM pre-training objectives.
Task
MAM
VAM
CAM
cheetah
14.9
46.2
64.5
finger
8.6
25.3
30.9
humanoid
39.4
135.6
243.5
quadruped
24.4
79.7
125.4
Table 8: Wall-clock pre-training time (minutes) for 100 epochs of pre-training.
Tables and g show that MAM yields the strongest or comparable downstream policy performance
across all environments and is consistently more efficient to train than both VAM and CAM.
19

---

Published as a conference paper at ICLR 2026
D ROBUSTNESS TO OFFLINE RL TRAINING OBJECTIVE
To test whether SPIN's benefits extend beyond the IQL objective used in the main experiments, we evaluate it with two additional offline RL methods
-enigge-Weigheed Actor-Critic (AWAC) (Nair et al. 2020) and Batch-Constrained Q-Learning (BCQ) (Fujimoto et al. 2019). These experi- ments are designed assess whether SPIN can serve as general framework for accelerating and improving offline RL in structured action spaces. For these experiments, we use the challenging and realistic medi i um-expert datasets.
D.I AWAC
Task
F-AWAC
AR-AWAC
SAINT
SPIN
cheetah
655.6 + 33.4
672.5 19.5 24.
675.6 17.1
finger
+ 1.6
2.9 3.7
681.8 + 339.5
863.1 6.9
humanoid
593.4 93.6
084.0
690.5
691.7 32.6
quadruped
813.8 + 50.3
801.6 + 41.3
821.5 + 44.0
826.4 + 35.7
Average
516.0
537.7
716.5
764.2
Table 9: Mean + std performance on DM Control tasks with AWAC variants. SPIN matches or exceeds the performance of the best baseline in every environment.
Task
F-AWAC
AR-AWAC
SAINT
SPIN
cheetah
68.6
125.0
122.2
97.6
finger
0.5
0.7
1.7
1.1
humanoid
84.5
82.0
72.8
3.3
quadruped
39.6
59.0
34.3
14.6
Total Runtime
193.2
266.7
231.0
116.6
Table 10: Wall-clock training time (minutes) to reach 95% of Factored AWAC asymptotic performance, with SPIN times including ASM pre-training.
Tables 9 and 10 report final performance and training efficiency for all architectures under the AWAC objective. SPIN achieves the highest average return (764.2), about 10% higher than the next-best baseline (SAINT). also the most efficient, attaining target performance in I 11..6 minutes, roughly twice fast SAINT.
20

---

Published as a conference paper at ICLR 2026
D.2 BCQ
Task
F-BCQ
AR-BCQ
SAINT
SPIN
cheetah
655.2
573.9
681.3 21.2
679.2 + 23.0
finger
090.0 69.9
49,66
23.
850.9 + 19.5
humanoid
597.2 + 45.1
698.4
43.6
712.4 + 31.7
quadruped
816.6 + 47.3
674.8 = 115.8
829.0 = 33.7
852.0 = 30.0
Average
689.7
651.9
760.8
773.1
Table 11: Mean + std performance on DM Control tasks with BCQ variants. SPIN matches or exceeds the performance of the best baseline in every environment.
Task
F-BCQ
AR-BCQ
SAINT
SPIN
cheetah
29.1
33.4
30.6
18.3
finger
3.8
1.0
1.6
0.3
humanoid
80.8
81.1
31.4
4.0
quadruped 40.8
43.7
17.0
11.6
Total Runtime
154.5
159.2
80.6
34.2
Table 12: Wall-clock training time (minutes) to reach 95% of Factored BCQ asymptotic performance, with SPIN times including ASM pre-training.
Tables 11 and 12 show that the trend holds with BCQ the base algorithm. SPIN attains the highest average performance (773.1). Its efficiency advantage also pronounced SPIN reaches target performance in 34.2 minutes, more than twice fast SAINT and over four times faster than the factored BCQ variants.
21

---

Published as a conference paper at ICLR 2026
E PERFORMANCE ON MAZE
To assess generalization beyond the DM Control locomotion suite, we evaluated all methods on configurable Maze navigation task introduced Beeson (2024). Unlike locomotion, this environment emphasizes goal-reaching rather than complex dynamics. Action dimensionality was varied by changing the number of actuators from 5 15.
Num. Actuators
F-IQL
AR-IQL
SAINT
SPIN
98.2 + 2.1
99.4 0.0
98.3 + 2.2
99.4 + 0.1 99.4 0.0
10
99.4 + 0.0
98.4 2.2
99.4
15
99.4 + 0.0
99.4 + 0.0
99.4 + 0.0
99.4 + 0.0
Average
99.0
99.1
99.0
99.4
Table 13: Performance across actuator configurations on the Maze environment. All methods achicve near- expert performance, indicating the task is solvable by all architectural variants.
Num. Actuators
F-IQL
AR-IQL
SAINT
SPIN
5
4.5
2.6
4.1
1.1
10
3.9
3.6
6.5
3.0
15
8.1
11.6
12.1
7.0
Total Runtime
16.5
17.8
22.7
11.1
Table 14: Wall-clock training time (minutes) reach F-IQL asymptotic performance, with SPIN times including ASM pre-training.
Table 13 shows that all methods, including factored and autoregressive baselines, achieve near-expert performance on the Maze task. This suggests that coordination between sub-actions is less demanding than in locomotion tasks and can be captured even without strong sub-action coordination.
Despite similar asymptotic returns, SPIN retains consistent efficiency advantage. As shown in Table SPIN reaches the F-IQL performance target fastest across all actuator configurations. With total runtime of 11.1 minutes across experiments, SPIN about 1.5x faster than the quickquickest baseline (F-IQL).
22

---

Published as a conference paper at ICLR 2026
F COMPARISON TO PRAJECTORY-CENTRIE PRE-TRAINING
To validate that SPIN's effectiveness comes from its action-centric masked modeling rather than pre-training itself, we compare its performance that of REPREM (Cai al. 2023). REPREM also follows a pre-train-finetune paradigm, but like most existing RL pre-training methods (see Section 2) its objective is trajectory- action sequences. Because RePreM's per-sample cost scales with horizon length, training on the 2 x 10 -ta-trannitoon medi um expert datasets prohibitively expensive. We therefore conduct the comparison on the random-medi um-expert datasets with x 10 transitions, which remain heterogeneous and challenging.
predicting future states and rewards from interleaved
Return
Runtime
Task
RePreM
SPIN
RePreM
SPIN
cheetah
272.5 + 27.1
332.4 + 43.1
585.8
1.8
finger
706.9 40.0
773. 2 + 14.9
599.3
1.0
humanoid
180.6 + 23.9
330.2 2 28.6
580.8
4.4
quadruped
328.8 95.3
561.0 + 70.6
585.7
2.9
Average
372.2
499.2
587.9
2.5
Table 15:
Performance and pre-training runtime on the random medi expeex datasets. SPIN achieves higher returns with over x00 faster training.
Results in Table 15 show that SPIN achieves average return of 499.2, while REPREM reaches only 372.2, demonstrating that in combinatorial action spaces, modeling the action manifold provides a stronger representation for control than trajectory prediction. The efficiency gap is even larger. SPIN completes training in 10.1 minutes, whereas REPREM requires 2351.6 6 minutes, more than x slower. This difference reflects the Morrizon-indeennce design of the ASM.
23


Reviews:
[{"review_id": "H3w8fNx8rB", "reviewer": "Reviewer_Mkwm", "summary": "The authors propose an RL method to handle large discrete action spaces. The method pretrains a transformer with masked action inputs to reconstruct the action, then do RL on top of the learned transformer representation. They show this method matches other large action RL baselines in just a fraction of their training times. The method is evaluated in a modified version of DM Control where the action space is discretized.", "questions": "See weaknesses, I would like to see more realistic experiments. For method novelty, it would be addressed by better framing, and also comparing against a standard representation learning approach like masked reconstruction over the entire input sequence.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The method is very simple and intuitive. By doing self supervised learning on the large action space, one can learn a more meaningful action representation than the original one. This will lead to large downstream gains. \n- The paper is well written and easy to understand.\n- The experiment section has interesting analysis results to pin down why SPIN is helpful.", "weaknesses": "### Empirical evaluation feels toy and contrived\n- these methods are all evaluated in rather artificial RL tasks (hopper, quadruped, etc.), where they take a popular benchmark (DM Control) and then factorize the action space. While useful for fast iteration and initial scientific insight, it is insufficient for convincing me that this method, or even the problem of large discrete action space, is useful. The authors motivated the problem by citing natural problems with large action spaces like recommender systems, robot assembly, etc. Could the authors show results in a more realistic problem setting?\n\n\n### Method novelty \n- In terms of methodological novelty, there's not too much at the high level. When you have noisy or high dimensional data, e.g. noisy sensors, high dim images, representation learning is the first thing we try to improve the signal to noise ratio of our data. So doing this for actions, using a standard masked reconstruction objective, seems very obvious, and not too \"novel\". On the other hand, this method is \"novel\" in the sense of applying the masked reconstruction objective to this particular problem where action spaces are large. \n- This can be seen as a special case of literature studying masked transformers for decision making problems [1], where the mask of the transformer is just set to the action modality. It would be interesting to compare SPIN against a masked transformer baseline that does masking over both state and action modalities. \n\n[1] Masked Trajectory Models for Prediction, Representation, and Control\n\n[2] PASTA: Pretrained Action-State Transformer Agents", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r76Akjcm3n", "reviewer": "Reviewer_98NV", "summary": "The authors proposed an algorithm for RL with combinatorial action spaces. The proposed method has two stages. We learn the action structure during the first stage, and then we learn  a policy in the second stage.", "questions": "Could the authors provide an approach to choose the stopping time of pretraining?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "By separating the learning of action structure and policy, the proposed algorithm overcomes the computational cost issue that a previous work named SAINT has.", "weaknesses": "Determining when to finish pretraining and move on to policy training is crucial. Stopping pretraining too early could lead to poor action structure modeling (Sec 6.1 illustrates the importance of sufficient pretraining), and stopping pretraining too late could lead to the same computational cost issue that is with SAINT. The authors do not provide an approach to choose the stopping time of pretraining. \n\nThe proposed method largely reuses the policy architecture in SAINT, and thus the novelty of this work is limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P0u9CIch50", "reviewer": "Reviewer_rC8Q", "summary": "This paper introduces SPIN, a two-stage approach designed to improve efficiency in discrete combinatorial action spaces. Specifically, it separates representation learning from policy learning. In the first stage, an action structure model learns a representation function that captures the manifold of valid actions. In the second stage, this representation is frozen and reused, with lightweight policy heads built on top of the pre-trained action structure model. The experimental results demonstrate clear benefits in terms of both performance and efficiency.", "questions": "1. Please justify the experimental comparison. \n2. Please clarify the contributions with respect to the earlier works, especially, MERLION.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written and well motivated.\n\nThe proposed idea is straightforward, and the algorithm is compatible with actor–critic frameworks, which enhances its applicability across a wide range of settings. The experimental results demonstrate the superiority of the proposed approach compared with the three selected baselines.", "weaknesses": "While the focus on offline RL is relevant, it is not sufficiently justified in the paper. In particular, SAINT is originally an online approach, which has been used here in an offline setting for comparison. In my view, it is not entirely fair to claim that SAINT jointly learns the action structure and control, as it was designed for a different purpose. This raises questions about the validity of the comparison.\n\nThe evaluation is also somewhat limited. The implementation details for the selected baselines are not described clearly, making the fairness of the comparison uncertain. While it is understandable that the authors aimed to keep architectural choices consistent, comparing directly with the original implementations of the baseline methods would strengthen the credibility of the results.\n\nThere are a few relevant works in this area that the authors may wish to consider for experimental comparison, such asOHIO (https://openreview.net/forum?id=dTPz4rEDok), and MERLION (https://proceedings.mlr.press/v162/gu22b/gu22b.pdf). Particularly, the paper claims the decoupling the representation learning from control. However, MERLION also learns reusable action embeddings. The contribution over MERLION remains unclear in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "em5GMnqAAJ", "reviewer": "Reviewer_UAbS", "summary": "Problem: \n\nThe paper tackles offline reinforcement learning in large, discrete combinatorial action spaces, settings where the agent must select from exponentially many joint actions (composed of multiple sub-actions) and ensure these selected sub-actions form coherent combinations. This is relevant for domains like healthcare decision support, robotics, recommendation systems, and fleet management, where online exploration is costly, risky, or infeasible.\n\nApproach:\n- The authors propose Structured Policy Initialization (SPIN), a two-stage framework that decouples representation learning from control. :\n(a) Action Structure Model (ASM) is trained to learn an action representation function, \n(b) ASM is frozen and lightweight policy heads are trained for downstream RL control on this learned action representation.", "questions": "- The paper notes that SPIN is compatible with IQL and AWAC but not CQL. Could you elaborate on stability issues that arise with value-regularization methods and whether hybrid objectives could reconcile them?\n\n- Why was masked conditional modeling chosen over contrastive or next-sub-action prediction? Did you test alternative pretext tasks, and if so, how did they compare?\n\n- Have you evaluated SPIN on higher-arity combinatorial domains (e.g., VRP or job-shop scheduling) where the sub-action semantics differ? Would the same ASM formulation apply without state–action token alignment?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- SPIN offers a principled and empirically validated way to accelerate and improve offline RL in large discrete combinatorial action spaces, primarily by decoupling structure learning and control, thus making learning tractable and robust as complexity grows. The separation of structure and control is motivated and clearly shown to overcome the slowness/instability of joint learning \n\n- SPIN works with multiple offline RL algorithms (IQL, AWAC, BCQ). Overall, SPIN is a promising and elegant approach that reframes discrete combinatorial control as a representation problem.", "weaknesses": "- The current framework requires architectural compatibility between ASM and policy modules for effective weight transfer. This can limit its integration with arbitrary RL architectures and restrict broader applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)