Structured Policy Initialization (SPIN)

• Structured Policy Initialization (SPIN) is a two-stage method that decouples structure learning from control policy optimization to address challenges in trajectory optimization and RL with combinatorial actions.
• It first learns valid action or trajectory manifolds via techniques like masked modeling or sequential convex programming, then fine-tunes lightweight controllers over these structured embeddings.
• Empirical results show SPIN significantly improves sample efficiency and performance, with up to +39% return gains in RL benchmarks and near-perfect initialization in complex control tasks.

Structured Policy Initialization (SPIN) denotes a two-stage methodological framework that systematically decouples the learning of structured action or trajectory representations from control policy optimization. It originated to address the limitations found in both trajectory optimization for nonlinear control and reinforcement learning (RL) with combinatorial action spaces. By first acquiring explicit or implicit knowledge of valid or feasible action/trajectory structure—either through guided search with local linearizations or masked conditional modeling—and then leveraging this structure for subsequent policy optimization, SPIN improves both sample efficiency and downstream performance across complex control and RL domains (Kim et al., 2021, Landers et al., 7 Jan 2026).

1. Motivation and Problem Domains

Structured Policy Initialization arises in two distinct but fundamentally related settings:

• Trajectory Optimization: Many trajectory optimization approaches for nonlinear systems are highly sensitive to initial guesses owing to nonconvex constraints and nonlinear dynamics. Poor initialization slows convergence or yields infeasible solutions.
• Reinforcement Learning with Combinatorial Actions: Offline RL agents with large, discrete, composite action spaces ($A = A_1 \times ... \times A_N$, $|A|=m^N$) face the curse of dimensionality and coherence: naïvely treating each sub-action independently ignores dependencies, while joint learning is often computationally infeasible or unstable.

SPIN addresses these challenges by explicitly learning action or trajectory manifolds before fine-tuning lightweight controllers, mitigating the intractable search over invalid or suboptimal spaces and providing high-quality initializations (Kim et al., 2021, Landers et al., 7 Jan 2026).

2. Stage 1: Explicit Structure Learning

SPIN starts by isolating and learning the structure of valid actions or trajectories, distinct from control policies.

2.1 Action Structure Model (ASM) in RL Contexts

In the offline RL context, an Action Structure Model (ASM) is trained via a permutation-equivariant Transformer encoder $f_\mathrm{ASM}$, taking as input state ($M$ tokens) and $N$ sub-action tokens. Training employs a masked modeling loss:

$$\mathcal{L}_{ASM} = \mathbb{E}_{(s,a)\sim \mathcal{D}} \; \mathbb{E}_{M\subset \{1..N\}} \sum_{i\in M} \mathrm{CrossEntropy}(f_i(h_a^i), a^i)$$

where masked slots are reconstructed via slot-specific heads. This learns a state-conditioned manifold of coherent actions, sharply reducing invalid or incoherent combinations at policy optimization time (Landers et al., 7 Jan 2026).

2.2 Trajectory Iterates via SCP in Trajectory Optimization

For trajectory optimization, structure learning consists of generating iterates by local linearization and a single step of Sequential Convex Programming (SCP) about previous reference trajectories $(\bar{x}_k^i, \bar{u}_k^i)$. The associated convex subproblem is:

$$\begin{align*} \min_{x^i_k, u^i_k, \nu^i_k} \quad & \sum_{i=1}^N \sum_{k=0}^{K-1} \left[\ell(t_k, x^i_k, u^i_k) + w_\nu \|\nu^i_k\|_1 + w_{trp} \|u^i_k - \pi_\theta(\bar{x}^i_k)\|_2^2\right] \ \text{s.t.}\quad & x_{k+1}^i = A_k^i x_k^i + B_k^i u_k^i + z_k^i + \nu_k^i \ & x_k^i \in \mathcal{X}(t_k), \quad u_k^i \in \mathcal{U}(t_k) \ & C_k^i x_k^i + D_k^i u_k^i + r_k^i \leq 0 \ & x_0^i = x_{\mathrm{init}}^i, \quad x_{K-1}^i = x_{\mathrm{final}}^i \end{align*}$$

Here $(\hat{x}_k^i, \hat{u}_k^i, \hat{\nu}_k^i)$ serve as structured, locally feasible trajectory candidates (Kim et al., 2021).

3. Stage 2: Policy Optimization over Structured Embeddings

The second phase in SPIN leverages the learned structure for efficient and effective policy optimization.

3.1 RL Policy Head over Frozen ASM

SPIN "freezes" the trained ASM encoder and reuses it as the front end for policy optimization. Lightweight policy heads (MLPs $g_i$) are attached to each sub-action embedding, yielding per-slot logits $|A|=m^N$0:

$|A|=m^N$1

This factorization, while structurally simple at the output, admits complex cross-dimensional dependencies through the ASM backbone (Landers et al., 7 Jan 2026). Updates are performed by any compatible actor-critic offline RL algorithm:

$|A|=m^N$2

where $|A|=m^N$3 is an advantage weighting dependent on the RL algorithm (e.g., IQL, AWAC, BCQ).

3.2 Supervised Policy Update in Trajectory Optimization

In the trajectory context, supervised learning is performed by minimizing:

$|A|=m^N$4

with data $|A|=m^N$5 generated via LTV-LQR feedback perturbations around the SCP iterates. This yields policies able to instantly propose near-feasible, near-optimal initial trajectories (Kim et al., 2021).

4. Algorithmic Workflow and Pseudocode

The core procedure in both settings is a two-stage, alternating training/optimization procedure:

Stage	RL Context (Landers et al., 7 Jan 2026)	Trajectory Context (Kim et al., 2021)
1. Structure Learning	Masked modeling of ASM transformer (Algorithm 1)	Trajectory update via one-step SCP
2. Policy Optimization	Freeze ASM, attach policy heads, RL fine-tuning (Algorithm 2)	Supervised update of NN policy

Both approaches explicitly decouple structure from control: ASM or SCP generates a representation/manifold of feasible action/trajectory structure, then control heads or NN policies are trained over these representations for efficient fine-tuning. The original publications provide explicit pseudocode for both stages.

5. Empirical Evaluation and Analysis

5.1 Reinforcement Learning Benchmarks

SPIN outperforms factorized (F-IQL) and fully joint Transformer (SAINT) policies on discretized DeepMind Control Suite tasks, improving average return by up to +39%, and reducing time to convergence by up to $|A|=m^N$6 in the most combinatorial tasks (Landers et al., 7 Jan 2026). Scaling analysis shows performance and efficiency gains increase with action dimension and granularity.

Further ablations show the bulk of SPIN's benefit arises from representation learning: even distilling a simple MLP policy to mimic frozen ASM embeddings matches Transformer policy performance, running $|A|=m^N$7 faster than SAINT. Early adaptation rates (after 10K steps) reach $|A|=m^N$8 of F-IQL's final return, compared to $|A|=m^N$9-$f_\mathrm{ASM}$0 for baselines. Representation quality, as assessed by linear probes, correlates with downstream RL success.

5.2 Trajectory Optimization Results

For a 6-DoF powered descent guidance problem (rocket landing), SPIN-trained policies yield dramatically better trajectory initializations: $f_\mathrm{ASM}$1 success rate and $f_\mathrm{ASM}$2 mean PTR iterations to convergence (vs. $f_\mathrm{ASM}$3 and $f_\mathrm{ASM}$4 for straight-line initialization). Maximum constraint violation is reduced from $f_\mathrm{ASM}$5 to $f_\mathrm{ASM}$6, with $f_\mathrm{ASM}$7 smaller errors on final conditions (Kim et al., 2021).

6. Implementation Notes and Robustness

• RL Implementations: PyTorch 2.6, NVIDIA A40 GPU, 5 seeds per experiment.
• Dataset Scope: Offline RL datasets span $f_\mathrm{ASM}$8 to $f_\mathrm{ASM}$9 transitions, four standard data regimes (Random, Medium, Expert, and Mixed).
• Policy Flexibility: SPIN's policy phase is compatible with IQL, AWAC, and BCQ (but not currently with value-regularized methods like CQL due to intractable global Q maximization).
• Robustness to Objective: SPIN consistently outperforms or matches SOTA RL methods regardless of the specific actor-critic update (see original appendix experiments for AWAC/BCQ).

7. Scope, Limitations, and Future Directions

Structured Policy Initialization reframes both RL with combinatorial actions and trajectory optimization as representation learning problems. By separating structure capture (ASM, SCP) from control, SPIN achieves scalability and convergence unattainable with monolithic end-to-end learning in high-dimensional spaces (Kim et al., 2021, Landers et al., 7 Jan 2026). Limitations include incompatibility with value regularization schemes requiring explicit Q-maximization over $M$0; handling of ordered or sequential sub-actions beyond permutation invariance; and generalization in sparse or biased data regimes. Addressing these remains an open area for future algorithmic and theoretical advances.