Technical Summary: Policy Optimization in Hybrid Discrete-Continuous Action Spaces via Mixed Gradients

Problem Statement

The paper addresses a fundamental limitation in Reinforcement Learning (RL) for large-scale problems: the credit assignment problem in high-dimensional action spaces. Standard model-free policy gradient methods, such as REINFORCE and Proximal Policy Optimization (PPO), rely on score-function (SF) estimators. These methods estimate gradients by correlating policy perturbations with cumulative rewards. In high-dimensional settings, this approach suffers from severe signal-to-noise degradation, as a single scalar return must be attributed to thousands of action components entangled with exogenous randomness.

While differentiable simulation (pathwise or PW gradients) offers a solution by backpropagating through the simulator to obtain sharper credit assignment, it traditionally requires fully smooth dynamics. Many practical problems in robotics, control, and operations involve hybrid discrete-continuous action spaces (e.g., selecting a discrete regime and optimizing a continuous control within it) or piecewise-smooth dynamics with discontinuities. In these settings, standard PW gradients become biased or uninformative, forcing practitioners to revert to less efficient SF-based methods and discard the available structural information.

The core challenge is to develop a policy optimization framework that leverages the smoothness of continuous components via backpropagation while correctly handling discrete decisions and discontinuities without introducing bias.

Methodology: Hybrid Policy Optimization (HPO)

The authors propose Hybrid Policy Optimization (HPO), a framework designed for exogenous MDPs where a simulator has access to the full realization of exogenous randomness (e.g., demand, weather) during training, allowing for counterfactual policy evaluation on the same trajectory.

1. Policy Representation

The policy $\pi$ is factorized into two components:

• Discrete Component ($\pi^X_\phi$): A stochastic policy that selects a discrete action $x_t$ (e.g., a regime or mode) from a finite set $\mathcal{X}$.
• Continuous Component ($\pi^B_\kappa$): A deterministic controller that outputs a candidate continuous action $b_t$ for each possible discrete action $j \in \mathcal{X}$. The executed action is $b_t = \pi^B_\kappa(s_t)[x_t]$.

This "towered" structure allows the continuous controller to be smooth with respect to its parameters $\kappa$, conditional on the discrete action.

2. The Mixed Gradient Estimator

The central theoretical contribution is an unbiased mixed gradient estimator that combines pathwise (PW) and score-function (SF) terms. For a policy parameterized by $\theta = (\phi, \kappa)$, the gradients are:

• Discrete Gradient ($\nabla_\phi J$): Computed using the standard SF estimator:
  $$\nabla_\phi J(\theta) = \mathbb{E}_\theta \left[ \sum_{t=0}^\infty \gamma^t \nabla_\phi \log \pi^X_\phi(x_t | s_t(\kappa)) Q^{\pi_\theta}(s_t(\kappa), a_t) \right]$$
• Continuous Gradient ($\nabla_\kappa J$): Computed as the sum of a Pathwise term and a Cross term:
  $$\nabla_\kappa J(\theta) = \underbrace{\mathbb{E}_\theta \left[ \sum_{t=0}^\infty \gamma^t \nabla_\kappa c(s_t(\kappa), (x_t, b_t(\kappa)), \xi_t) \right]}_{\text{PW Gradient Term}} + \underbrace{\mathbb{E}_\theta \left[ \sum_{t=0}^\infty \gamma^t \nabla_\kappa \log \pi^X_\phi(x_t | s_t(\kappa)) Q^{\pi_\theta}(s_t(\kappa), a_t) \right]}_{\text{Cross Term}}$$

Key Mechanism:

• The PW term backpropagates through the simulator for the continuous actions, holding the discrete actions fixed. This isolates the direct effect of continuous actions on costs, significantly improving credit assignment.
• The Cross term captures how changes in the continuous parameters $\kappa$ alter the state trajectory $s_t(\kappa)$, which in turn shifts the probability distribution of the discrete policy $\pi^X_\phi$.
• The estimator is implemented using stop_grad (detach) operations in autodifferentiation frameworks (e.g., PyTorch) to ensure the discrete sampling is treated as fixed for the PW term, while the cross term retains the dependency of the discrete log-probability on the state.

3. Handling Discontinuities

The paper demonstrates that problems with action discontinuities (Piecewise-Smooth MDPs) can be reformulated as hybrid MDPs. By introducing a discrete action to select the "smooth region" of the action space and a continuous action to specify the control within that region, the framework preserves unbiasedness while exploiting local smoothness.

Key Contributions

1. HPO Framework: A policy optimization framework for hybrid discrete-continuous action spaces that yields an unbiased mixed gradient estimator. It leverages PW gradients wherever the problem structure permits, avoiding the brittleness of pure PW methods and the high variance of pure SF methods.
2. Theoretical Characterization of the Cross Term: The authors prove that the Cross term vanishes as the discrete policy approaches a best response (Theorem 2). Specifically, the magnitude of the cross term is bounded by the suboptimality of the discrete policy ($\epsilon$) and the sensitivity of the state to continuous parameters. This result implies that near optimality, the coupling between discrete and continuous components weakens, enabling approximate decentralized updates.
3. Empirical Validation: Extensive experiments on Joint Replenishment Problems (JRP) and Switched Linear Quadratic Regulators (S-LQR) demonstrate that HPO substantially outperforms PPO, particularly as the continuous action dimension increases.
   • Scaling: HPO's performance gap over PPO widens significantly with the dimension of the continuous action space ($p$). For example, in S-LQR, PPO fails to converge for $p > 24$, while HPO converges reliably up to $p=48$.
   • Gradient Quality: Analysis shows that the mixed estimator maintains high alignment and low RMSE in high dimensions, whereas pure SF estimators (used by PPO) collapse to near-zero alignment as dimensionality grows.
4. Variance Reduction Strategy: The paper shows that dropping the cross term (HPONoCross) can be beneficial near optimality. Since the cross term's expectation vanishes while its variance grows as the discrete policy improves, removing it reduces noise without introducing significant bias in the final stages of training.

Results and Significance

The paper claims that HPO represents a significant step toward applying policy-gradient methods to realistic, large-scale problems where differentiable simulators are available but discrete decisions have historically limited the reach of pathwise methods.

• Performance: In the JRP and S-LQR benchmarks, HPO consistently achieves lower test losses and requires fewer policy updates to reach target performance gaps compared to PPO. The advantage is most pronounced in high-dimensional continuous spaces ($p \ge 40$).
• Gradient Efficiency: The mixed gradient provides a drastically better signal than pure SF gradients. In high dimensions, the SF estimator's alignment with the true gradient direction drops to near zero, while the mixed estimator maintains alignment above 0.8.
• Decentralization Potential: The theoretical result regarding the vanishing cross term suggests a pathway toward semi-decentralized training, where discrete and continuous components can be optimized separately once the discrete policy is sufficiently close to a best response.

The authors position this work as a bridge between model-free RL and differentiable simulation, extending the latter's applicability to hybrid action spaces without sacrificing the unbiasedness required for convergence. The code and resources are made available to facilitate further research in this direction.