XiNet: Online Training & Pruning in Deep RL

• XiNet is a framework that enables simultaneous online training and pruning in deep reinforcement learning by incorporating stochastic Bernoulli gates.
• It employs DenseNet-style OFENet feature extractors and cost-aware regularization to balance performance with computational efficiency.
• XiNet optimizes both network weights and gate parameters in a unified training loop, resulting in compact models that outperform smaller networks trained from scratch.

XiNet is a framework for online training and pruning in deep reinforcement learning that integrates stochastic unit-wise gating into reinforcement learning networks enhanced by the Online Feature Extractor Network (OFENet). It was introduced to address the compute and memory growth that accompanies scaling neural networks in reinforcement learning, particularly when DenseNet-style feature extractors are used to improve representation quality and downstream policy and value learning. Rather than applying pruning post hoc, XiNet jointly optimizes the network weights and the parameters of variational Bernoulli gates that can deactivate units during training, yielding compact deterministic networks at convergence (Guenter et al., 16 Jul 2025).

1. Problem setting and conceptual scope

The immediate setting for XiNet is deep reinforcement learning with OFENet-enhanced agents. OFENet augments reinforcement learning by learning high-dimensional feature representations $z$ from observations and state-action pairs, which can stabilize training and support larger policy and value networks. The same scaling, however, increases parameter counts and FLOPs substantially because DenseNet-style concatenation propagates earlier features into later layers, and the enlarged downstream reinforcement learning networks inherit the increased feature dimensionality. XiNet is designed to preserve the advantages of starting from large models while reducing their eventual computational footprint (Guenter et al., 16 Jul 2025).

The central objective is simultaneous training and pruning. Let $\theta$ denote trainable network weights and $\phi$ denote the parameters of variational Bernoulli gates. XiNet optimizes

$$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$

where $L_{\mathrm{RL}}$ encapsulates the reinforcement learning losses under a sampled gate realization $\xi$, and $R(\phi)$ is a sparsity-inducing regularizer. The stochastic formulation is intended to drive the variational parameters of low-utility units toward zero; once a gate parameter falls below a pruning threshold, the corresponding structure is rendered permanently inactive and removed.

Within this formulation, XiNet is not a compression method appended after training. Its defining property is that sparse structure discovery is part of the learning problem itself. This distinction matters because the reported results compare online-pruned large networks not only to unpruned large networks, but also to smaller networks trained from scratch.

2. Architecture and gating mechanism

XiNet is defined over OFENet-enhanced reinforcement learning architectures. OFENet consists of two DenseNet-style multilayer perceptrons: a state feature extractor $\phi_o$, which maps observations $o_t$ to features $z_{o_t}=\phi_o(o_t)$, and a state-action feature extractor $\theta$0, which maps $\theta$1 to features $\theta$2. DenseNet-style connectivity concatenates the gated output of each block with a pass-through of its previous input. This redundancy is useful for scaling, but it also magnifies downstream cost (Guenter et al., 16 Jul 2025).

XiNet introduces binary random gates at unit granularity. Each gate vector $\theta$3 multiplies a layer’s activation vector element-wise. For the observation pathway,

$\theta$4

and the block output concatenates $\theta$5 with a pass-through term $\theta$6, where $\theta$7 is diagonal. The state-action pathway is analogous, with

$\theta$8

The use of diagonal pass-through scalers $\theta$9 in place of a literal identity is intended to improve pruning efficiency. The activation functions satisfy $\phi$0, and batch normalization is part of each dense block.

The gates are variational Bernoulli random variables:

$\phi$1

The distribution factorizes across layers and units, and the optimized quantities are the per-unit Bernoulli parameters $\phi$2. XiNet applies the same unit-wise gating principle not only inside OFENet but also at each hidden layer of the policy, value, and critic networks.

3. Stochastic objective, auxiliary training, and regularization

XiNet adopts a stochastic optimization perspective in which both weights and gate parameters are optimized. In the generic form reused per batch in reinforcement learning, the objective is

$\phi$3

where

$\phi$4

Under the flattening hyperprior, the gate regularizer is

$\phi$5

with $\phi$6. Its derivative is constant in $\phi$7, so it provides a flat tunable push toward zero.

A key structural identity is the linearity of $\phi$8 in each gate parameter:

$\phi$9

Here $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$0 and $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$1 are the expected losses conditioned on $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$2 or $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$3. The difference $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$4 is the unit’s marginal “usefulness” in the loss.

For the reinforcement learning component, XiNet uses Soft Actor-Critic (SAC). The actor loss is

$$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$5

with $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$6 in sampling mode during training and $$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$7 in evaluation mode. The critic losses are

$$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$8

where

$$J(\theta,\phi)=\mathbb{E}_{\xi\sim q(\xi;\phi)}\!\left[L_{\mathrm{RL}}(\theta,\xi)\right]+R(\phi),$$9

The temperature loss is

$L_{\mathrm{RL}}$0

If a value network is used, its loss is also defined explicitly in the formulation.

OFENet is trained concurrently through a one-step next-state prediction auxiliary task:

$L_{\mathrm{RL}}$1

The corresponding auxiliary objective is

$L_{\mathrm{RL}}$2

For the reinforcement learning networks $L_{\mathrm{RL}}$3, XiNet adds

$L_{\mathrm{RL}}$4

and the total reinforcement learning loss is

$L_{\mathrm{RL}}$5

The overall design therefore couples policy optimization, representation learning, and model compression within a single training process.

4. Cost-aware sparsity and expected complexity

A distinctive feature of XiNet is that the sparsity regularization is matched to expected computational cost, rather than being an undifferentiated penalty on gate activations. Because the gates are Bernoulli random variables, the expected active width of a layer is the $L_{\mathrm{RL}}$6 norm of its gate-parameter vector. This makes it possible to write explicit expected-complexity expressions for OFENet, its auxiliary predictor, and the downstream policy and value networks (Guenter et al., 16 Jul 2025).

For the observation feature extractor $L_{\mathrm{RL}}$7, with observation dimension $L_{\mathrm{RL}}$8,

$L_{\mathrm{RL}}$9

For the state-action extractor $\xi$0, with action dimension $\xi$1,

$\xi$2

The $\xi$3 terms account for batch normalization and bias operations, and the $\xi$4 terms account for the diagonal pass-through $\xi$5.

The predictor cost is

$\xi$6

For any downstream reinforcement learning network $\xi$7, the first-layer input dimension depends on the OFENet output width, and the corresponding input-stage cost is

$\xi$8

XiNet then aggregates these terms into a weighted OFENet complexity:

$\xi$9

where $R(\phi)$0 discounts training-only costs. Setting $R(\phi)$1 emphasizes deployment compactness, $R(\phi)$2 emphasizes training-time savings, and intermediate values interpolate between those objectives.

The cost-aware scheme chooses flattening hyperparameters so that the gate regularizer matches expected complexity:

$R(\phi)$3

and for each feed-forward reinforcement learning network $R(\phi)$4,

$R(\phi)$5

This construction is meant to remove much of the manual per-layer hyperparameter tuning that would otherwise accompany sparsity-inducing priors. A plausible implication is that XiNet’s regularization is not only sparsity-promoting but also topology-sensitive, since it penalizes units in proportion to their effect on the larger DenseNet-style computation graph.

5. Optimization procedure and pruning dynamics

Gate parameters are updated with a straight-through (ST) estimator. At each training step, XiNet samples gate realizations for the gated layers, performs forward and backward propagation on the instantiated subnetwork, updates the weights by standard backpropagation, and then updates the gate parameters by projected gradient descent:

$R(\phi)$6

The marginal usefulness term is approximated by

$R(\phi)$7

For the flattening regularizer, $R(\phi)$8; for cost-matched regularization $R(\phi)$9, the derivative becomes $\phi_o$0 (Guenter et al., 16 Jul 2025).

The ST estimator is described as biased but low variance and computationally efficient. Concrete or Gumbel-Softmax relaxations and REINFORCE are identified as possible alternatives, but they were not required in the reported experiments. The ST approximation is also given an analytical interpretation as a first-order Taylor approximation.

The training loop interleaves data collection, OFENet updates, SAC updates, and pruning. During data collection, the agent computes $\phi_o$1 and samples $\phi_o$2 in evaluation mode, using $\phi_o$3 rather than sampled $\phi_o$4 to avoid additional exploration noise. Transitions are stored in a replay buffer. OFENet updates then sample minibatches, sample $\phi_o$5 and $\phi_o$6, optimize $\phi_o$7, update $\phi_o$8, $\phi_o$9, and $o_t$0, and project $o_t$1 back to $o_t$2. SAC updates sample gates for the network currently being updated, augment its loss with the corresponding regularizer, update its weights and gates, and update target Q networks as in SAC.

Pruning is triggered when a gate parameter satisfies $o_t$3. Units may be pruned immediately or after a brief grace period, after which their incident weights are removed from the computational graph. In practice, the gate parameters are often initialized to 1 and held fixed during an initial warmup of 200k steps. During the final 20% of training, any non-converged gate parameters are rounded to $o_t$4 or $o_t$5, producing deterministic pruned networks. Because gating occurs before DenseNet concatenation and because pass-through connections are reweighted by diagonal $o_t$6, zeroed gates remove a unit’s future contribution without invalidating the concatenation structure.

The paper further states that, with multilinear cost-matched regularization and the linear-in-$o_t$7 objective structure, the gate subproblem’s minimizer lies at a vertex of the hypercube $o_t$8, so $o_t$9 in supervised settings. In reinforcement learning, where the data distribution is nonstationary, the same deterministic convergence is reported empirically rather than as a general theorem.

6. Experimental results, interpretation, and limitations

The reported evaluation uses MuJoCo continuous-control tasks via Gymnasium: HalfCheetah-v2, Ant-v2, Hopper-v2, and Walker2D-v2. The reinforcement learning agent is Soft Actor-Critic. Baseline SAC uses two hidden layers with 256 units; SAC-Big uses two hidden layers with 1024 units. OFE uses $z_{o_t}=\phi_o(o_t)$0 and $z_{o_t}=\phi_o(o_t)$1 DenseNets with 6 layers for Ant, Hopper, and Walker2D, 8 layers for HalfCheetah, and 32 units per OFE layer, together with SAC networks of width 256. OFE-Big increases OFE layers to 128 units and SAC networks to width 512. The pruning experiments start from OFE-Big scale and prune both OFENets and SAC networks online; three regularization settings, Prune-A, Prune-B, and Prune-C, vary $z_{o_t}=\phi_o(o_t)$2, $z_{o_t}=\phi_o(o_t)$3, and $z_{o_t}=\phi_o(o_t)$4 to explore performance-complexity trade-offs (Guenter et al., 16 Jul 2025).

Training uses Adam with learning rates $z_{o_t}=\phi_o(o_t)$5 and $z_{o_t}=\phi_o(o_t)$6. Gate parameters are initialized to 1 and held for 200k steps; in the final 20% of training, they are rounded to $z_{o_t}=\phi_o(o_t)$7. Total environment steps are $z_{o_t}=\phi_o(o_t)$8 for Ant, Hopper, and Walker2D, and $z_{o_t}=\phi_o(o_t)$9 for HalfCheetah. The replay buffer is warm-started with $\theta$00 random-policy steps, followed by short OFENet pretraining on those samples.

On HalfCheetah-v2, the reported return and deployment/train parameter counts are: SAC 12757 with 74k/289k; SAC-Big 15534 with 1081k/4300k; OFE 16609 with 174k/821k; OFE-Big 17299 with 1286k/6276k; Prune-A 16774 with 152k/1466k; Prune-B 16033 with 28k/390k; and Prune-C 13593 with 21k/236k. On Ant-v2, SAC achieves 4974 with 100k/389k, SAC-Big 6345 with 1182k/4695k, OFE 7161 with 187k/799k, OFE-Big 7608 with 1057k/5114k, Prune-A 7229 with 242k/1111k, Prune-B 6517 with 73k/576k, and Prune-C 5829 with 41k/266k.

On Hopper-v2, SAC achieves 3197 with 71k/280k, SAC-Big 3534 with 1069k/4264k, OFE 3548 with 139k/651k, OFE-Big 3699 with 924k/4573k, Prune-A 3659 with 55k/556k, Prune-B 3515 with 15k/246k, and Prune-C 3470 with 8k/102k. On Walker2D-v2, SAC achieves 3998 with 74k/289k, SAC-Big 4976 with 1081k/4301k, OFE 5468 with 143k/663k, OFE-Big 5490 with 934k/4613k, Prune-A 5679 with 361k/2187k, Prune-B 5152 with 67k/586k, and Prune-C 4803 with 50k/389k.

The reported interpretation is that starting large and pruning online produces compact agents with minimal performance loss, and that in several tasks Prune-B and Prune-C remain close to OFE or OFE-Big while using far fewer parameters for both deployment and training. A common misconception is that pruning a large network during training is equivalent to training a small network from scratch; the reported comparisons do not support that equivalence, since pruned models often achieve higher returns at similar or lower parameter counts than baseline SAC, with Hopper-v2 and Walker2D-v2 given as examples.

The paper also reports structure-level pruning patterns. The OFENet state extractor $\theta$01 is pruned heavily across tasks, whereas $\theta$02 exhibits task-dependent patterns: later layers remain wider in Ant, earlier layers remain wider in Walker2D, and Cheetah is more uniform. This suggests that the cost-aware regularization allocates sparsity nonuniformly across the DenseNet hierarchy where it most reduces downstream cost.

Several limitations are identified. Performance-sparsity trade-offs depend on $\theta$03, and $\theta$04; overly aggressive regularization can over-prune and reduce returns, while weak regularization leaves large models. Stochastic gating early in training introduces additional noise that can interact with SAC exploration and occasionally cause evaluation spikes, with Cheetah cited as an example. Early OFENet layers are particularly delicate because they influence many concatenations and downstream input dimensions, so non-cost-aware pruning can become imbalanced. Proposed extensions include convolutional layers, other reinforcement learning algorithms such as TD3 and PPO, additional auxiliary tasks, adaptive $\theta$05, and combinations of unit-level with block-level or layer-level gates.