Primal-Dual Representation Learning for Low-Rank Constrained MDPs

Chenhao Zhou; Chao Zhang; Wensong Bai; Hanbin Zhao; Hui Qian

cs.LGOct 22, 2026

Primal-Dual Representation Learning for Low-Rank Constrained MDPs

Authors: Chenhao Zhou, Chao Zhang, Wensong Bai, Hanbin Zhao, Hui Qian

Organizations: College of Computer Science and Technology, Zhejiang University

Abstract

Representation learning has substantially improved the efficiency of unconstrained continuous-state Markov Decision Processes (MDPs), but its extension to Constrained MDPs (CMDPs) is difficult because the transition representation must be learned while constraint violation is controlled. In this paper, we study representation learning for low-rank CMDPs with continuous state spaces, and propose REP-PD-TEG for the soft cumulative-violation setting. In the feature learning period, the algorithm interacts with the environment using a decaying- $ε$ -greedy mixture execution policy, and uses the collected samples to learn transition features by maximum likelihood. Based on the learned features, it constructs uncertainty bonuses and plans with an optimistic Lagrangian objective to obtain the return policy. We prove regret and soft-violation guarantees for the return policies, and quantify the influence of the decaying- $ε$ -greedy exploration. In particular, REP-PD-TEG achieves the optimal $\widetilde O(\sqrt K)$ return policy bound, and gives a $\widetilde O(K^{3/4})$ overall sample bound when exploratory actions are charged. For the stricter hard cumulative violation, where feasibility must be verified episode by episode, we propose REP-PD-H-TEG. It plans under a bonus-corrected lower-confidence utility certificate with the learned representation, and performs a stabilized dual search to obtain the return policy, yielding regret and hard-violation bounds. Together, these results give the first theoretically guaranteed representation-learning method for low-rank continuous-state CMDPs with both soft cumulative guarantees and certified hard-violation control.

Figures & tables


MDP	Method	Known Feature	Regret	Violation
Tabular	Efroni et al. [16]	Yes	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{3}\|\mathcal{A}\|H^{4}K}\right)$	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{3}\|\mathcal{A}\|H^{4}K}\right)$
Liu et al. [27]	Yes	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{3}\|\mathcal{A}\|H^{6}K}\right)$	0

Linear	Ding et al. [13]	Yes	$\tilde{O}\left(\sqrt{d^{2}H^{6}K}\right)$	$\tilde{O}\left(\sqrt{d^{2}H^{6}K}\right)$
Linear	Ghosh et al. [17]	Yes	$\tilde{O}\left(\sqrt{d^{3}H^{3}K}\right)$	$\tilde{O}\left(\sqrt{d^{3}H^{3}K}\right)$

Table 1: Comparison of soft-violation bounds. The two REP-PD-TEG rows report the complexity of the best return policy specialization and the best execution policy specialization, respectively.


MDP	Method	Known Feature	Known Safe Action Set	Regret	Violation
Tabular	Efroni et al. [16]	Yes	Yes	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{3}\|\mathcal{A}\|H^{4}K}\right)$	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{3}\|\mathcal{A}\|H^{4}K}\right)$
Bura et al. [9]	Yes	Yes	$\tilde{O}\left(\sqrt{\|\mathcal{S}\|^{2}\|\mathcal{A}\|H^{6}K}\right)$	0

Linear	Amani et al. [5]	Yes	Yes	$\tilde{O}\left(\sqrt{d^{3}H^{4}K}\right)$	0
	Kitamura et al. [25]	Yes	Yes	$\tilde{O}\left(\sqrt{d^{5}H^{8}K}\right)$	0

Table 2: Comparison of hard-violation bounds. The two REP-PD-H-TEG rows report the complexity of the best return policy specialization and the balanced execution policy specialization, respectively.

\widehat{P}_{k}=(\widehat{\mu}_{k},\widehat{\phi}_{k})\in\arg\max_{(\mu,\phi)\in\mathcal{M}}\mathbb{E}_{D_{k}+D^{\prime}_{k}+D^{\prime\prime}_{k}}\big[\log\mu(x^{\prime})^{\top}\phi(x,u)\big].

Algorithm 1 REP-PD-TEG

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

\widehat{P}_{k}=(\widehat{\mu}_{k},\widehat{\phi}_{k})\in\arg\max_{(\mu,\phi)\in\mathcal{M}}\mathbb{E}_{D_{k}+D^{\prime}_{k}+D^{\prime\prime}_{k}}[\log\mu(x^{\prime})^{\top}\phi(x,u)].

Appendix

Algorithm 2 REP-PD-H-TEG with an explicit one-sided certificate

Figure 1: Illustrative single-run comparison of REP-PD-TEG, REP-UCB, and the known-feature linear method of Ghosh et al. [17] on the job-scheduling CMDP.

Environment	Raw Ret.	Raw Cost	REP-PD-NN Ret.	REP-PD-NN Cost
PointGoal1	$-1.38\pm 0.58$	$27.85\pm 18.99$	$\mathbf{-0.14\pm 0.03}$	$\mathbf{16.45\pm 12.69}$
PointPush1	$-1.40\pm 0.60$	$32.17\pm 14.66$	$\mathbf{-0.17\pm 0.12}$	$\mathbf{12.13\pm 6.95}$
PointPush2	$-1.40\pm 0.71$	$13.00\pm 8.17$	$\mathbf{-0.03\pm 0.01}$	$\mathbf{10.90\pm 3.85}$

Appendix

Table 3: Safety Gym results for the practical neural implementation REP-PD-NN . Metrics are the mean and standard deviation over

3

seeds at the final deterministic evaluation checkpoint after

1

M environment steps.

Experiment	Worker	Steps	Runs	Time/run	Total compute
Job scheduling	CPU ( $<1$ GB RAM)	–	1	$<1$ min	$<1$ CPU-hour
Safety Gym PointGoal1	1 RTX 4090 (24GB)	$10^{6}$	6	8–20 h	48–120 GPU-hours
Safety Gym PointPush1	1 RTX 4090 (24GB)	$10^{6}$	6	8–20 h	48–120 GPU-hours
Safety Gym PointPush2	1 RTX 4090 (24GB)	$10^{6}$	6	8–20 h	48–120 GPU-hours

Appendix

Table 4: Approximate compute required to reproduce the reported experiments. Runtime denotes wall-clock time for one run. Total compute is over the reported methods and seeds.

Markov Decision Processes

Markov decision processes model sequential choices under uncertainty, linking actions, state transitions, and rewards to planning and reinforcement learning.

See all 157 papers

Feb 16, 2026

Near-Optimal Sample Complexity for Online Constrained MDPs

Safety is a fundamental challenge in reinforcement learning (RL), particularly in real-world applications such as autonomous driving, robotics, and healthcare. To address this, Constrained Markov Decision Processes (CMDPs) are commonly used to enforce safety constraints while optimizing performance. However, existing methods often suffer from significant safety violations or require a high sample complexity to generate near-optimal policies. We address two settings: relaxed feasibility, where small violations are allowed, and strict feasibility, where no violation is allowed. We propose a model-based primal-dual algorithm that balances regret and bounded constraint violations, drawing on techniques from online RL and constrained optimization. For relaxed feasibility, we prove that our algorithm returns an

\varepsilon

-optimal policy with

\varepsilon

-bounded violation with arbitrarily high probability, requiring

\tilde{O}\left(\frac{SAH^3}{\varepsilon^2}\right)

learning episodes, matching the lower bound for unconstrained MDPs. For strict feasibility, we prove that our algorithm returns an

\varepsilon

-optimal policy with zero violation with arbitrarily high probability, requiring

\tilde{O}\left(\frac{SAH^5}{\varepsilon^2ζ^2}\right)

learning episodes, where

ζ

is the problem-dependent Slater constant characterizing the size of the feasible region. This result matches the lower bound for learning CMDPs with access to a generative model. Our results demonstrate that learning CMDPs in an online setting is as easy as learning with a generative model and is no more challenging than learning unconstrained MDPs when small violations are allowed.

Chang Liu, Yunfan Li, Lin F. Yang

University of California, Los Angeles

Constrained RL Safe RL Sample-Efficient RL Model-Based RL

Ask PDF

Oct 1, 2026

Rate-Optimal Algorithm for Adversarial Linear CMDPs

We study episodic adversarial linear constrained Markov decision processes (CMDPs) with unknown transitions, where both the loss and constraint functions may vary adversarially across episodes. The best previous algorithm achieves

\widetilde{\mathcal{O}}(K^{3/4})

regret and cumulative constraint violation, leaving a gap to the optimal

\widetilde{\mathcal{O}}(\sqrt{K})

dependence on the number of episodes

K

. We close this gap by proposing a new primal dual algorithm that achieves

\widetilde{\mathcal{O}}(\sqrt{K})

regret and cumulative constraint violation without assuming Slater's condition. We further extend the algorithm to achieve the same

\widetilde{\mathcal{O}}(\sqrt{K})

guarantees for regret and hard constraint violation, which does not allow constraint violations to cancel across episodes. The main challenge is that learning linear CMDPs requires uniform concentration over a value function class with a controlled covering number, whereas standard techniques in constrained online learning, such as policy mixing, can make this class more complex. Our algorithm combines adaptive Follow the Regularized Leader (FTRL), contracted value estimation, and an exponential Lyapunov function. An adaptive dual regularizer offsets the dependence on the dual weights in the primal regret bound, removing the need for policy mixing. We further show that the normalization in the FTRL update bounds the policy parameters independently of the magnitudes of the dual weights, which explains why the resulting policy class remains compatible with uniform concentration. Under feature access, the computational complexity is independent of the size of the state space.

Kihyun Yu, Honghao Wei, Dabeen Lee

KAIST · Washington State University · Seoul National University

Regret Minimization in RL Markov Decision Processes Constrained RL Linear Markov Decision Processes

Ask PDF

Mar 29, 2026

Near-Optimal Primal-Dual Algorithm for Learning Linear Mixture CMDPs with Adversarial Rewards

We study safe reinforcement learning in finite-horizon linear mixture constrained Markov decision processes (CMDPs) with adversarial rewards under full-information feedback and an unknown transition kernel. We propose a primal-dual policy optimization algorithm that achieves regret and constraint violation bounds of

\widetilde{O}(\sqrt{d^2 H^3 K})

under mild conditions, where

d

is the feature dimension,

H

is the horizon, and

K

is the number of episodes. To the best of our knowledge, this is the first provably efficient algorithm for linear mixture CMDPs with adversarial rewards. In particular, our regret bound is near-optimal, matching the known minimax lower bound up to logarithmic factors. The key idea is to introduce a regularized dual update that enables a drift-based analysis. This step is essential, as strong duality-based analysis cannot be directly applied when reward functions change across episodes. In addition, we extend weighted ridge regression-based parameter estimation to the constrained setting, allowing us to construct tighter confidence intervals that are crucial for deriving the near-optimal regret bound.

Kihyun Yu, Seoungbin Bae, Dabeen Lee

Department of Industrial and Systems Engineering, KAIST · Department of Mathematical Sciences, Seoul National University · Research Institute of Mathematics, Seoul National University +2

Adversarial MARL Reinforcement Learning Constrained RL Markov Decision Processes Regret Minimization in RL

Ask PDF

Abstract

Markov Decision Processes

Explore similar work

Near-Optimal Sample Complexity for Online Constrained MDPs

Rate-Optimal Algorithm for Adversarial Linear CMDPs

Near-Optimal Primal-Dual Algorithm for Learning Linear Mixture CMDPs with Adversarial Rewards