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

Near-Optimal Sample Complexity for Online Constrained MDPs
Feb 16, 2026Chang Liu, Yunfan Li, Lin F. YangConstrained RL Safe RL
University of California, Los Angeles
Rate-Optimal Algorithm for Adversarial Linear CMDPs
Oct 1, 2026Kihyun Yu, Honghao Wei, Dabeen LeeRegret Minimization in RL Markov Decision Processes
KAIST · Washington State University · Seoul National University
Near-Optimal Primal-Dual Algorithm for Learning Linear Mixture CMDPs with Adversarial Rewards
Mar 29, 2026Kihyun Yu, Seoungbin Bae, Dabeen LeeAdversarial MARL Reinforcement Learning
Department of Industrial and Systems Engineering, KAIST · Department of Mathematical Sciences, Seoul National University · Research Institute of Mathematics, Seoul National University +2

Feb 16, 2026J/K move · Enter open · S save

Near-Optimal Sample Complexity for Online Constrained MDPs

Chang Liu, Yunfan Li, Lin F. Yang

University of California, Los Angeles

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

Open paper AskPDF arXiv

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.

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