PaperlayerView exact revision on arXiv

Source-linked AI summary

Primal-Dual Representation Learning for Low-Rank Constrained MDPs

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

arXiv:2610.30756v1cs.LG

TL;DR

CMDP representation learning must discover unknown transition features while controlling constraint violation, a setting not covered by known-feature methods. The paper develops primal-dual methods for soft and hard cumulative violations, with return- and execution-policy guarantees whose rates depend on probing.

  • Problem

    CMDPs with unknown transition features remain insufficiently addressed because representation learning must support reward optimization while constraint violation is controlled.

    Source: §1 · p. 1

  • Method

    The methods learn transition features from ϵ-greedy mixture-policy data by MLE, then use learned-feature bonuses in optimistic Lagrangian planning for soft violation or a lower-confidence utility certificate and stabilized dual search for hard violation.

    Source: §1 · p. 2, §3.1 · p. 6

  • Results

    REP-PD-TEG achieves optimal ~O(K^1/2) return-policy regret and soft violation at p=0; p=1/4 gives ~O(K^3/4) execution-policy regret and soft violation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The work provides a theoretically guaranteed representation-learning method for low-rank continuous-state CMDPs with soft cumulative guarantees and hard-violation control.

    Source: §5 · p. 10

    The guarantees concern cumulative violation control and do not imply per-episode safety during learning.

    Source: §2 · p. 5

Abstract

from arXiv · show

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-$\epsilon$-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-$\epsilon$-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.

1 Introduction

The paper addresses CMDPs with unknown transition features, where representation learning must support reward optimization while cumulative constraint violation is controlled. It introduces separate methods for soft and hard violation criteria and analyzes how probing affects return- and execution-policy guarantees.

1 Introduction

  • The soft-violation method learns transition features by MLE from ϵ-greedy probing, then plans with learned-feature bonuses in an optimistic Lagrangian objective.The Lagrangian bonus-augments both reward and dual-weighted utility terms.Source: §1 · p. 2
  • REP-PD-TEG achieves optimal ~O(K^1/2) return-policy regret and soft violation at p=0; p=1/4 yields ~O(K^3/4) execution-policy regret and soft violation.The execution-policy bound includes the O(K^(1-p)) complexity added by ϵ-greedy exploration.Source: §1 · p. 2
  • REP-PD-H-TEG obtains return-policy regret ~O(K^(3/4+p)) and hard violation ~O(K^(1/2+p)); its p=1/8 execution specialization has ~O(K^7/8) regret and hard violation.The return-policy specialization uses p=0 and has ~O(K^1/2) hard-violation dependence.Source: §1 · p. 2
  • For hard violation, REP-PD-H-TEG uses a bonus-corrected lower-confidence utility certificate and stabilized dual search to obtain its return policy.Hard violation charges each episode’s shortfall separately, unlike aggregate soft violation.Source: §1 · p. 2

1.1 Related Work

  • CMDP methods with known low-dimensional features do not address the unknown-feature setting, while prior low-rank CMDP work studies near-optimal feasible-policy identification after training rather than guarantees over the learning trajectory.Representation error in CMDPs affects reward optimization and constraint evaluation, coupling feature coverage with safety-aware policy optimization.Source: §1 · pp. 1–2, §1.1 · p. 3
  • Prior representation-learning work in low-rank MDPs does not address online CMDP safety constraints; this paper presents a representation-learning framework for that setting.The cited prior methods include model-based and model-free approaches as well as reward-free exploration and multi-agent Markov games.Source: §1.1 · p. 4

2 Background

The background formulates an episodic CMDP as reward maximization subject to an expected-utility threshold, then defines regret and cumulative soft and hard violation measures. Its guarantees use a low-rank transition model and a finite realizable function class in an occupancy-based stochastic-oracle setting.

2 Background

  • The learning objective maximizes expected reward subject to an expected-utility threshold b, with the paper focusing on one constraint while noting that the method extends to multiple constraints.The setup allows constraint violations during learning and measures regret and total violations over K episodes.Source: §2 · pp. 4–5
  • Soft violation measures cumulative signed shortfall, whereas hard violation sums only positive shortfalls; neither metric guarantees per-episode safety during learning.Zero soft violation may still accompany hard violation that grows linearly with episodes.Source: §2 · p. 5
  • The transition model is assumed to have rank d, with transition probabilities represented by inner products of state-action and state embeddings.The paper assumes normalized embeddings; the low-rank structure includes tabular, linear, and block CMDPs as special cases.Source: §2 · pp. 5–6
  • The stochastic oracle samples occupancy-distributed states and auxiliary transitions; K counts oracle rounds, not raw environment transitions, and bounds do not charge complete roll-in trajectories.The expected transition count for K queries is O(K/(1 − γ) + 3K); execution-policy corollaries concern the stationary ϵ-mixture, not the entire probing trajectory.Source: §2 · p. 6
  • Formal guarantees assume a finite realizable model class containing the true transition model and access to an MLE oracle for selecting a transition estimate.This restricts candidate representations rather than state-space cardinality.Source: §2 · p. 6

3 Soft Violation via Cumulative Primal-Dual Control

REP-PD-TEG addresses soft cumulative violation by learning transition features from exploratory data and using uncertainty-aware optimistic primal-dual planning. Its polynomial probing schedule trades return-policy guarantees against the cost of executing exploratory actions.

3 Soft Violation via Cumulative Primal-Dual Control

  • Soft violation aggregates constraint deficits and surpluses across episodes before taking the positive part, enabling utility optimism in the primal-dual update.The return-policy sequence is evaluated separately from auxiliary tuples used for model estimation and bonus construction.Source: §3 · p. 6, §3.1 · p. 7

3.1 REP-PD-TEG Algorithm

  • REP-PD-TEG uses a decaying-ϵ mixture of the previous return policy and uniform actions to collect data for transition-feature MLE.The schedule ϵ_k = k^−p makes the exploration–deployment tradeoff explicit; return-policy terms scale as eO(K^1/2+p).Source: §3.1 · p. 6, §3.2 · p. 8
  • The learned features support a merged-covariance uncertainty bonus that augments both reward and utility estimates for optimistic planning.The projected dual variable is bounded in [0, ξ], with ξ = 2/θ in the analysis.Source: §3.1 · pp. 6–7

3.2 Polynomial-Schedule Guarantees

  • The schedule tradeoff makes p = 0 statistically sharp for return policies, while p = 1/4 balances representation learning with deployed exploration cost.At p = 1/4, the execution-policy gap is O(K^3/4/(1 − γ)).Source: §3.2 · p. 8

4 Hard Violation via Per-Episode Certification

REP-PD-H-TEG addresses hard violation by requiring per-episode feasibility certification rather than relying on utility optimism. It retains the representation-learning procedure while changing policy selection to use lower-confidence certificates and a capped dual search.

4.1 Lower-Certificate Planning in REP-PD-H-TEG

  • REP-PD-H-TEG preserves tertiary data collection, MLE representation updates, merged covariance, and bonus construction, but certifies each deployed policy with a lower-confidence utility bound.The first candidate whose certificate meets b is deployed; if none passes before the cap, the final capped-dual candidate is used.Source: §4 · p. 8, §4.1 · pp. 8–9

4.2 Hard-Violation Guarantees

  • Hard violation is sublinear for p < 1/2, whereas regret is sublinear only for p < 1/4.The K^1/2+p hard-violation term reflects representation learning, while an additional K^1/4 in regret comes from capped in-episode dual search.Source: §4.2 · p. 10
  • For p = 1/8, execution-policy regret and hard violation are both eO(K^7/8), up to problem-dependent factors.The exploratory-action cost determines this balanced execution-policy rate.Source: §4.2 · p. 10

5 Conclusion

The paper presents two representation-learning algorithms for low-rank continuous-state CMDPs, with guarantees for soft and hard cumulative violation. The conclusion also records theoretical assumptions and the scope of the appendix analyses.

5 Conclusion

  • REP-PD-TEG and REP-PD-H-TEG provide return- and execution-policy guarantees for soft- and hard-violation control, respectively.The hard method uses a lower-confidence utility certificate and stabilized dual search.Source: §5 · p. 10
  • The theorems assume finite realizability and exact MLE and planning; the hard result additionally assumes a positive optimistic-score gap.The bounds do not charge full roll-in trajectories, and the experiments do not exhaustively evaluate hard certification.Source: §5 · p. 10

Appendix Roadmap

  • The appendix roadmap assigns shared proof tools, hard-constraint ingredients, soft decaying-ϵ analysis, uniform-probing specialization, hard-certificate analysis, and experiments to Appendices A–H.Appendix B collects MLE, elliptical-potential, simulation, and primal-dual lemmas.Source: Appendix Roadmap · p. 14, B Auxiliary Lemmas · p. 14

A Notation and Auxiliary Concentration

  • The notation appendix recalls an empirical-to-population covariance comparison underlying the uncertainty bonus, first for uniform probing and later for tertiary decaying-ϵ designs.Its proof cites the standard comparison for uniform probing and applies the same argument to predictable auxiliary streams.Source: A Notation and Auxiliary Concentration · p. 14

C Hard-Constraint Technical Lemmas

  • The hard-constraint technical appendix uses lower-confidence certificates and Slater-induced control of cap episodes, rather than optimistic feasibility checks, as its hard certificate.Its positive-action-gap assumption separates the optimistic Lagrangian maximizer from suboptimal actions.Source: C Hard-Constraint Technical Lemmas · p. 15

D Decaying ϵ-Greedy Probing: Tertiary-Aware Analysis

The decaying-ϵ analysis repairs a coverage gap in the naive two-level probing rule by adding a tertiary probe level. Its guarantees concern return policies, with execution-policy transfer and violation control handled separately.

D Decaying ϵ-Greedy Probing: Tertiary-Aware Analysis

  • The naive two-level decaying-ϵ probing rule does not provide the coverage needed to close the optimism proof.The appendix identifies a minimal tertiary-aware repair that restores the missing denominator.Source: D Decaying ϵ-Greedy Probing:… · p. 16
  • The appendix proves a regret guarantee for the return-policy sequence under the repaired probing algorithm.The result follows the two-level obstruction and the addition of a tertiary probe level.Source: D Decaying ϵ-Greedy Probing:… · p. 16
  • Execution-policy transfer is established afterward; the same front end supports soft-violation control when cumulative optimistic utility overestimation is bounded.For hard violation, later analyses replace optimistic feasibility checks with a lower-confidence certificate.Source: D Decaying ϵ-Greedy Probing:… · p. 16

D.1 Why the Naive Two-Level Variant Is Not Enough

This appendix explains why the naive two-level probing rule fails to support the required coverage comparison and analyzes a tertiary-aware repair, including supporting lemmas for its guarantees. It also establishes that, under the nonbinding-cap condition, the mixture-bonus bound holds for the actual algorithmic bonus using simulation-lemma comparisons and bonus domination.

D.1 Why the Naive Two-Level Variant Is Not Enough

  • The naive two-level rule admits no universal pointwise comparison, motivating a third probe level for the merged-covariance design.A primary-only covariance alternative such as Uehara et al. [41] is not ruled out, but would require a different potential argument.Source: §D.1 · p. 16

D.2 The Tertiary-Aware Repaired Algorithm

  • The repaired algorithm adds a tertiary probe-action and transition pair while retaining bonus covariance merged over primary and secondary tuples.The appendix analyzes the main-text Algorithm 1; the change from the earlier two-stream protocol is in auxiliary sampling.Source: §D.2 · p. 17

D.3 Supporting Lemmas

  • The supporting lemmas provide simultaneous tuple-level MLE control and a pathwise-to-averaged bridge for the MLE-mismatch class.The bridge is used to transfer internal MLE-mismatch terms; generic test functions are handled through deterministic coverage coefficients.Source: §D.3 · pp. 18–19
  • Further lemmas establish merged-covariance comparison, tertiary coverage, a tertiary-aware one-step-back inequality, and direct optimism for the repaired bonus.The one-step-back and optimism statements apply under the events specified in their respective lemmas.Source: §D.3 · pp. 21–23, Also define · p. 24

Also define

  • The one-step-back bound applies to every policy and bounded measurable test function on the specified lemma events.The proof uses a pathwise-to-averaged transfer for the MLE mismatch, while generic test-function dependence is handled through deterministic coverage coefficients.Source: Also define · pp. 23–24
  • The proof compares the estimated transition with the optimal transition using a simulation lemma, then invokes Lemma D.8 and bonus monotonicity.It assumes normalized discounted reward lies in [0, 1] and uses pointwise domination of the population bonus by the empirical bonus.Source: Also define · p. 25
  • Under the nonbinding-cap condition, the actual algorithmic bonus inherits the same bound as the mixture empirical bonus.The passage states this transfer without providing the bound’s expression.Source: Also define · pp. 24–25

D.4 Final Aggregation

The final aggregation combines return-policy bounds with execution-policy transfer for decaying exploration, then specializes the analysis to uniform probing. The uniform case retains the repaired tertiary-aware design.

D.4 Final Aggregation

  • The theorem-level aggregation assumes the required high-probability events and a nonbinding cap.The proposition also specifies tuning for its regret and soft-violation dependence.Source: §D.4 · p. 29
  • For p ∈ [1/4, 1/2), execution-policy regret and soft violation match the asymptotic orders of the return-policy bounds.The execution term is no larger order than the return-policy terms in this range.Source: §D.4 · pp. 32–33
  • ˜O(K^(1/2+p)) bounds hold for return-policy regret and soft violation under the decaying schedule, with the additional execution term lower order when p > 0.The regret bound carries an extra θ^-1 dependence.Source: §D.4 · pp. 30, 31, 33

E Uniform Exploration as a Tertiary-Aware Repaired Variant

  • The uniform-exploration specialization sets ϵk ≡ 1, making all probe actions uniform while retaining the tertiary probe level and merged bonus covariance.It serves as a clean benchmark for the repaired analysis.Source: E Uniform Exploration as a… · p. 33

E.1 Uniform Specialization

  • The uniform specialization is Algorithm 1 with probing kernels fixed to U(A); its tertiary tuple, merged covariance, model fitting, and primal-dual update remain unchanged.Source: §E.1 · p. 34

E.2 Supporting Specializations

  • Under uniform probing, coverage coefficients collapse to explicit constants, and the tertiary coefficients specialize to Γk = |A|.Source: §E.2 · p. 34

E.3 Theorem-Level Derivation

The soft-violation analysis adds a summable control on optimistic utility overestimation to the representation-learning and optimism arguments. The hard-violation analysis instead distinguishes the algorithm’s one-sided certificate correction from a policy-dependent width used only in the proof.

E.3 Theorem-Level Derivation

  • The cumulative overestimation term is summable under the repaired tertiary design, allowing the projected dual argument to yield the stated soft-violation bound.For uniform probing (p = 0), the appendix gives an explicit return-policy theorem with the familiar Õ(|A|d^2√K) dependence, up to problem-dependent factors and logarithmic terms.Source: §E.3 · p. 35

F Hard Constraints Under Decaying Tertiary Probing

  • The repaired hard-violation analysis uses a one-sided correction in the algorithm and a two-sided width only in the proof, without an external certificate event or a uniformly summable radius over inner-loop policies.Source: F Hard Constraints Under Decaying… · p. 36

F.1 Algorithm and Certificate

  • The algorithmic correction ℓk is distinct from the policy-dependent width wk(π), which is not an algorithmic input.Source: §F.1 · p. 36

F.2 Derived Certificate Event

The hard-violation analysis builds a one-sided certificate and bounds its behavior for returned policies. It combines a hard front end, certificate-based regret reduction, and cap-policy violation analysis to establish rates under decaying exploration.

F.2 Derived Certificate Event

  • The one-sided certificate and selected-policy interval are established on the representation-learning event, subject to a nonbinding-cap condition.Source: §F.2 · p. 36

F.3 Hard Front End

  • The hard front end compares optimistic Lagrangian values at the selected dual level, with utility-side statistical terms scaled by Yk and softmax approximation adding value loss.Source: §F.3 · p. 39

F.4 Selected-Certificate Reduction

  • The selected-certificate reduction bounds regret using grid sensitivity and the certificate for the policy actually selected from the grid.Source: §F.4 · p. 39

F.5 Hard Violation at the Executed Cap Policy

  • When Ymaxθ ≥ 1, the cap-policy analysis bounds hard violation using Slater pressure on episodes that reach the cap; otherwise it gives a trivial bound.Source: §F.5 · p. 40

F.6 Polynomial Schedules

  • For ϵk = k^-p, p ∈ [0, 1/2), under Theorem F.3 and Proposition F.6 assumptions, Corollary F.7 gives YmaxKsgrid = eO(K1/2) with probability at least 1 −δ.Source: §F.6 · pp. 40–41

G The p = 0 Hard Specialization

The appendix specializes the derived-certificate analysis to uniform tertiary probing to make the return-policy rate explicit. Its hard-violation argument applies feasibility certification to break episodes and cap-level analysis to cap episodes.

G The p = 0 Hard Specialization

H.1 Job scheduling.

The job-scheduling experiment illustrates how learned representations, optimistic bonuses, and primal-dual updates interact under unknown dynamics. REP-PD-TEG trades higher regret than REP-UCB for substantially lower cumulative constraint violation in the reported run.

H.1 Job scheduling.

  • The toy CMDP models pending jobs over a horizon H = 10, with actions to withhold or submit jobs and rewards that make processing temporarily expensive.The state space is {0, 1, ..., 9}, where the state represents pending jobs.Source: §H.1 · p. 42
  • The implementation learns latent transition factors from collected transitions, updates the representation learner after each episode, then performs the primal-dual policy update with the induced exploration bonus.Source: §H.1 · p. 42
  • REP-PD-TEG has higher regret than REP-UCB but substantially lower cumulative constraint violation, and is comparable in scale to the known-feature linear method on both metrics.Figure 1 reports a single run without uncertainty bands, so the authors do not infer statistical significance or an empirical asymptotic rate.Source: §H.1 · p. 42

H.2 Practical Neural Implementation

REP-PD-NN adapts the primal-dual safe-RL structure for continuous-control tasks by replacing full-batch representation learning with a minibatch stochastic objective. On three reported Safety Gym tasks, it improves final evaluation return and reduces final evaluation cost relative to a raw-state baseline.

H.2 Practical Neural Implementation

  • REP-PD-NN replaces per-round full-batch MLE with a mini-batch stochastic representation objective, then uses learned representations in reward and cost critics and a Lagrangian-style policy update.The variant is evaluated on Safety Gym tasks.Source: §H.2 · p. 42
  • The raw-state baseline keeps the same Lagrangian training structure but removes representation updates and feeds raw state-action pairs to the actor and critics.This design attributes performance differences to the learned representation rather than a different safe-RL backbone.Source: §H.2 · p. 43
  • Across three Safety Gym tasks, REP-PD-NN improves final evaluation return and reduces final evaluation cost versus the raw baseline after 1M steps, averaging over 3 seeds.Table 3 reports the final deterministic evaluation checkpoint.Source: §H.2 · p. 43, Table 3 · p. 44

H.3 Computational Resources

  • Safety Gym’s 18 reported runs require approximately 144–360 GPU-hours total, or 18–45 wall-clock hours when eight GPUs run in parallel.Each run uses one GPU; preliminary sweeps and diagnostics used additional GPU time of the same order and were excluded from reported results.Source: §H.3 · p. 44

Related papers

Loading 2610.30756v1…