Environment Parameter Gradient Theorem for Policy-Environment Co-Design in Reinforcement Learning
arXiv:2607.12590v1 [eess.SY] 14 Jul 2026
Amber Srivastava Abstract— Reinforcement learning (RL) is traditionally concerned with learning a control policy for a fixed environment. In many engineering systems, however, the environment itself is alterable: physical or operational parameters can be tuned to shape the transition dynamics and costs experienced by the agent. This motivates jointly optimizing both the policy and the environment design parameters. To this end, we establish an Environment Parameter Gradient Theorem — a formal expression for the gradient of the value function with respect to environment parameters. The key theoretical device is a generalized action-value function Qπ,ξ (s, a, ζ), which comprises two copies of the environment parameters: ζ governs the cost and transition dynamics at the current state–action pair, while ξ governs the future rollouts. This decoupling yields a tractable closed-form gradient expression and is essential to the theorem’s derivation. Building on this result, we develop a model-free algorithm that simultaneously learns the optimal policy and the environment parameters. We demonstrate the efficacy of our framework on a UAV network design problem, where the optimal UAV placement (environment parameters) and communication routes (governed by the policy) are learned jointly to minimize the total communication cost in the network.
I. INTRODUCTION Reinforcement learning (RL) has emerged as a powerful framework for sequential decision-making under uncertainty, with applications spanning robotics, autonomous systems, and control [1]. In its standard formulation, RL focuses on adapting the agent policy while assuming that the underlying system configuration remains fixed. However, in many areas such as robotics [2], communication networks [3], and vehicle platooning [4], the system configuration itself is re-designable, and one can additionally optimize the configuration parameters along with the agent’s policy. These co-design problems have been addressed through several approaches in the literature. For instance, co-design methods in control theory typically assume a parametric model of the underlying system and jointly optimize the plant and controller in a model-based setting [5], [6]. More recently, RL-based methods have been proposed that maintain a probability distribution over candidate designs and adapt it alongside the policy. In particular, [7] uses Proximal Policy Optimization (PPO) for policy learning together with likelihood-ratio updates for adapting the design distribution, while [8] employs an off-policy actor–critic framework based This work was supported by the Prime Minister Early Career Research Grant from the Anusandhan National Research Foundation (ANRF/ECRG/2024/004876); and the Faculty Seed Grant Fund at IIT Delhi. Amber Srivastava is with Faculty of Mechanical Engineering, Indian Institute of Technology Delhi, India. email: [email protected]
on Deep Deterministic Policy Gradient (DDPG) together with entropy-regularized design-distribution updates. Population-based and evolutionary approaches have also been explored for joint robot morphology and control optimization. In [9], the design and policy parameters are jointly adapted using a population-based REINFORCE framework [10], while [11] combines evolutionary methods such as Genetic Algorithms and Bayesian Optimization with RL methods such as PPO for robot co-design tasks. In [2], the Soft Actor-Critic (SAC) algorithm [12] is employed with a critic explicitly conditioned on the design parameters, and design variables are subsequently extracted via particle swarm optimization. Other works [13], [14] combine policygradient methods with known differentiable environment dynamics to compute gradients of suitable loss functions with respect to system design parameters. More recent works have explored differentiable simulation, morphology-environment co-evolution, and bilevel formulations for co-design problems in robotics and embodied intelligence [15], [16]. Despite these advances, existing approaches primarily focus on algorithmic procedures for joint design and control, while comparatively little attention has been devoted to deriving explicit expressions for the sensitivity of the infinitehorizon value function with respect to environment design parameters. Such an expression — an analogue of the policy gradient theorem [17], [18] for environment parameters — would provide a principled foundation for (stochastic) gradient-based optimization of environment parameters in a model-free setting. Works that do compute gradients with respect to environment parameters [13], [14] rely on known differentiable model of the environment dynamics and do not address the model-free setting. To the best of our knowledge, no existing work derives a tractable expression for environment design parameters in infinite-horizon setting that is amenable to the model-free learning. The stochastic policy gradient (SPG) [10], [17] and the deterministic policy gradient (DPG) [18] theorems are landmark results that derive tractable gradient expressions for the value function with respect to policy parameters θ. Deriving such a result for environment design parameters ζ is difficult because ζ enters the objective through three channels simultaneously: the instantaneous cost, the transition dynamics, and the optimal policy all depend on ζ explicitly. Its dependence is therefore more deeply entangled than that of θ in the policy gradient setting [17], [18], where θ enters through a single channel, the action. The compounding dependencies of the objective J(π, ζ) on ζ give rise to two distinct challenges: deriving the gradient
term, and devising a model-free method to estimate it. The first is that, unlike in the policy gradient theorems, the gradient with respect to ζ cannot be written recursively in terms of successor states when the conventional actionvalue function is used. We restore this recursive structure by introducing a generalized action-value function Qπ,ξ (s, a, ζ) that carries two copies of the environment parameter, ζ and ξ, with distinct roles made precise in Section III. Briefly, ζ governs the instantaneous cost and transition dynamics at the current step from (s, a), while ξ governs the environment experienced in future rollouts under π. This decoupling yields a well-defined recursive sensitivity equation for the gradient with respect to environment parameter, from which a closed-form expression of the gradient follows. The second challenge is model-free estimability of the critic Qπ,ξ (s, a, ζ). The gradient theorem requires the critic to capture the sensitivity of the action-value function to the current-step environment parameter, which demands that the critic observe variation in ζ during training. This variation is obtained as follows: with the environment in state s, the design parameter is perturbed slightly from its baseline value ζ̄ before action a is applied, and the resulting instantaneous cost and successor state are used to train the critic, with future value estimates bootstrapped from a target critic operating at the original ζ̄. This creates a controlled variability between the environment parameter seen at the current step and the one used for future value estimation, directly operationalizing the two-copy separation between ζ and ξ required by the theorem. We emphasize that this separation is an analytical device required to compute ∇ζ J(π, ζ) (Section III); the final design is governed by a single environment parameter value, the same for both copies. This controlled variation in ζ can in principle be applied to a real physical system, but doing so could often be unsafe or operationally disruptive and sample-inefficient; we therefore perform it offline against a simulator, as is already standard for much of RL training. Our main contributions are as follows: (i) Environment Parameter Gradient Theorem: We establish a formal gradient theorem for environment parameters in infinite-horizon MDPs, providing a tractable closed-form expression for ∇ζ J(π, ζ) amenable to model-free estimation. To the best of our knowledge, this is the first such result in the literature. (ii) Generalized Action-Value Function: We introduce the two-copy action-value function Qπ,ξ (s, a, ζ) as the key theoretical device enabling the above theorem. (iii) Model-Free Algorithm: We develop a model-free actor–critic algorithm that jointly optimizes the policy and environment design parameters from sampled trajectories, without knowing the transition dynamics and the instaneous costs. We demonstrate the efficacy of the proposed framework on a UAV-based communication network design problem, where the algorithm co-optimizes UAV placement and packet routing policy. The resulting designs closely match (average objective value difference of 0.9%) those obtained when the transition dynamics and cost, along with their dependence on ζ, are explicitly known.
II. P ROBLEM F ORMULATION We consider a parameterized Markov decision process (MDP) M(ζ) = ⟨S, A, pζ , cζ , γ⟩, where S is a finite state space, A is a finite action space, pζ (s′ |s, a) is the transition kernel giving the probability of transitioning to s′ from s under action a and environment parameter ζ, cζ (s, a, s′ ) ≥ 0 is the instantaneous cost of the transition, and γ ∈ (0, 1] is the discount factor. The environment parameter ζ ∈ Z ⊆ Rd is a continuous design variable that directly governs both the transition dynamics and the instantaneous cost. Our results hold for discounted problems (γ < 1) and, under appropriate assumptions (such as existence of cost-free termination state and a proper policy) [19], for the undiscounted stochastic shortest path (SSP) setting (γ = 1) . At each time step t, the agent observes state st ∈ S, selects action at ∈ A according to a stochastic policy π(at |st , ζ), transitions to the next state st+1 ∼ pζ (·|st , at ), and incurs instantaneous cost cζ (st , at , st+1 ). The policy π(a|s, ζ) is explicitly conditioned on ζ, since the optimal action selection depends on the underlying MDP configuration. We denote by pπ,ζ the trajectory distribution induced by policy π and environment parameter ζ. The objective is to minimize the expected cumulative cost over both the policy π and the parameter ζ under an initial state distribution µ over S: X min J(π, ζ) := µ(s)Vπ,ζ (s), (1) π,ζ∈Z
s∈S
where the state-value function Vπ,ζ (s) is given by "∞ # X Vπ,ζ (s) = Epπ,ζ γ t cζ (st , at , st+1 ) s0 = s .
(2)
t=0
The above value function satisfies the Bellman recursion X Vπ,ζ (s) = π(a|s, ζ) Qπ,ζ (s, a), (3) a∈A
Qπ,ζ (s, a)=
X
pζ (s′ |s, a) cζ (s, a, s′ ) + γVπ,ζ (s′ ) , (4)
s′ ∈S
where Qπ,ζ (s, a) is referred to as the action-value function. For a fixed ζ, the minimization over π in (1) reduces to a standard RL problem [1]. The key challenge addressed in this work is the minimization over ζ, which requires the gradient of J(π, ζ) with respect to ζ. III. E NVIRONMENT PARAMETER G RADIENT T HEOREM Differentiating J(π, ζ) with respect to ζ through the standard action-value function Qπ,ζ (s, a) in (4) does not yield a tractable closed-form gradient: the resulting sensitivity term does not close into a recursion over successor states. This motivates a generalized action-value function that formally separates the environment parameter into two copies, each playing a distinct role. Definition 1: For a policy π, the generalized action-value function Qπ,ξ (s, a, ζ) is defined as: X Qπ,ξ (s, a, ζ)= pζ (s′ |s, a) cζ (s, a, s′ )+γVπ,ξ (s′ ) , (5) s′
where ζ parameterizes the current step transition kernel pζ (·|s, a) and instantaneous cost cζ (s, a, ·) at the (s, a) pair, while ξ parameterizes the environment over all subsequent steps, entering through Vπ,ξ (s′ ). We refer to ζ and ξ as the current and future step copies of the environment parameters, respectively. By treating the current-step and future environment parameters as independent variables, the generalized Q function isolates the sensitivity to the current-step parameter ζ from its effect on future rollouts, providing the structure needed to close the recursive sensitivity equation, as shown in the proof of Theorem 1. Note that when ζ = ξ, the generalized Q reduces to the standard action-value function Qπ,ζ (s, a) in (4). Further, the state-value function Vπ,ζ (s) in terms of the generalized action-value function is given by X Vπ,ζ (s) = π(a|s, ζ)Qπ,ζ (s, a, ζ). (6) a∈A
Theorem 1 (Environment Parameter Gradient (EnvPG)): Let π(·|s, ζ) be a differentiable stochastic policy and Qπ,ξ (s, a, ζ) be the generalized action-value function in (5). The gradient of the objective J(π, ζ) with respect to the environment parameter ζ, evaluated at ζ = ζ̄, is given by: " ∞ X ∇ζ J(π, ζ) = Epπ,ζ̄ γ t ∇ζ Qπ,ζ̄ (st , at , ζ)+ ζ=ζ̄
t=0
ln π(at |st , ζ) · Qπ,ζ̄ (st , at , ζ̄)
# ζ=ζ̄
s0 ∼ µ . (7)
Proof: Consider the definition of the state-value function Vπ,ζ (s, a) in terms of the generalized Q function in (6). Differentiating it with respect to ζ and evaluating at ζ = ζ̄: X ∇ζ Vπ,ζ (s) = ∇ζ π(a|s, ζ) ζ=ζ̄ Qπ,ζ̄ (s, a, ζ̄) ζ=ζ̄
a
+
X
π(a|s, ζ̄)∇ζ Qπ,ζ (s, a, ζ) ζ=ζ̄ . (8)
a
For the second term in (8), we expand ∇ζ Qπ,ζ (s, a, ζ) ζ=ζ̄ using its definition in (5) resulting into ∇ζ Qπ,ζ (s, a,ζ) ζ=ζ̄ = ∇ζ Qπ,ζ̄ (s, a, ζ) ζ=ζ̄ X +γ pζ̄ (s′ |s, a)∇ζ Vπ,ζ (s′ ) ζ=ζ̄ , (9) s′
where the first term captures the direct effect of ζ on the current-step cost and transition kernel with future parameters frozen at ζ̄, and the second term propagates the sensitivity recursively through the successor states. Substituting back: X ∇ζ Vπ,ζ (s) = ∇ζ π(a|s, ζ) ζ=ζ̄ Qπ,ζ̄ (s, a, ζ̄) ζ=ζ̄
a
+
X
π(a|s, ζ̄)∇ζ Qπ,ζ̄ (s, a, ζ) ζ=ζ̄
a
+γ
X a
π(a|s, ζ̄)
X
pζ̄ (s′ |s, a)∇ζ Vπ,ζ (s′ ) ζ=ζ̄ . (10)
s′
Applying the log-derivative trick to the first term, i.e., ∇ζ π(a|s, ζ) = π(a|s, ζ)∇ζ ln π(a|s, ζ), and combining the
first two terms we obtain h X = π(a|s, ζ̄) ∇ζ Qπ,ζ̄ (s, a, ζ) ∇ζ Vπ,ζ (s) ζ=ζ̄
a
+ ln π(a|s, ζ)Qπ,ζ̄ (s, a, ζ̄) ζ=ζ̄ i X ′ ′ . pζ̄ (s |s, a)∇ζ Vπ,ζ (s ) +γ ζ=ζ̄
s′
(11)
Unrolling the above recursion across time steps — which converges under discounting (γ < 1) or the SSP assumptions (γ = 1) [19] — and recognizing that the trajectory distribution under policy π and parameter ζ̄ is pπ,ζ̄ , we obtain " ∞ X = Epπ,ζ̄ γ t ∇ζ Qπ,ζ̄ (st , at , ζ)+ ∇ζ Vπ,ζ (s) ζ=ζ̄
t=0
ln π(at |st , ζ) · Qπ,ζ̄ (st , at , ζ̄)
# ζ=ζ̄
s0 = s . (12)
The gradient of J(π, ζ) in (7) follows immediately by taking the expectation of ∇ζ Vπ,ζ (s) ζ=ζ̄ over s0 ∼ µ. The gradient expression in Theorem 1 consists of two structurally distinct terms, each capturing a different channel through which ζ influences the objective. The first term, ∇ζ Qπ,ζ̄ (st , at , ζ) ζ=ζ̄ , captures the direct sensitivity of the action-value function to ζ through the instantaneous cost and transition dynamics at the current t-th step, with the future parameter frozen at ζ̄. The second term, ∇ζ ln π(at |st , ζ) · Qπ,ζ̄ (st , at , ζ̄), is a policy-gradient term: it weights the sensitivity of the policy to ζ by the action-value function, mirroring the score-function form of the standard policy gradient [17]. Both terms are estimable without a model of the environment: the first is obtained by differentiating a critic network that estimates Qπ,ξ (s, a, ζ), and the second from the score of an actor network that represents the policy π(a|s, ζ) together with the critic’s value; both critic and actor can be learned from sampled trajectories (see Section V). IV. S OLUTION M ETHODOLOGY A. Entropy Regularization and Gibbs Policy To encourage exploration and improve training, we augment the objective (1) with a common entropy regularization term [12], [20], yielding the entropy-regularized objective: JH (π, ζ)=Epπ,ζ
∞ hX
i γ t ct +α ln π(at |st , ζ) s0 ∼µ , (13)
t=0
where ct = cζ (st , at , st+1 ), and α > 0 is a temperature parameter. Applying the EnvPG Theorem 1 to the regularized objective (13) yields the following gradient expression: " ∞ X ∇ζ JH (π, ζ) = Epπ,ζ̄ γ t ∇ζ QH π,ζ̄ (st , at , ζ)+ ζ=ζ̄
t=0
#
ln π(at |st , ζ) · QH π,ζ̄ (st , at , ζ̄)
ζ=ζ̄
s0 ∼ µ ,
(14)
where the generalized entropy-regularized action-value function QH π,ξ (s, a, ζ) is defined analogously to the Definition 1: X ′ H ′ pζ (s′ |s, a) cH QH ζ (s, a, s )+γVπ,ξ (s ) , (15) π,ξ (s, a, ζ)= s′ ′ ′ where cH ζ (s, a, s ) = cζ (s, a, s ) + α ln π(a|s, ζ), ζ governs
the environment parameters in the current step, ξ governs the H future rollouts, and Vπ,ξ (s′ ) denotes the entropy regularized ′ value function from s with parameter ξ. Let πθ denote a policy parameterized by θ. The gradient of JH (πθ , ζ) with respect to θ follows from the standard entropy-regularized policy gradient [12]: ∞ hX ∇θ JH (πθ , ζ) = Epπθ ,ζ γ t QH πθ ,ζ (st , at , ζ) + α t=0
i ∇θ log πθ (at |st , ζ) s0 ∼ µ , (16) where the additive α arises from the entropy term in (13). We observe that under direct parameterization of the policy (a separate policy parameter for each state–action pair), the gradient expression of Theorem 1 reduces to a single term. Here, the optimal policy for the entropy-regularized objective (13) is known to be the Gibbs distribution [3], [20]: π ∗ (a|s, ζ) = P
exp{− α1 QH π ∗ ,ζ (s, a, ζ)}
, 1 H ′ a′ ∈A exp{− α Qπ ∗ ,ζ (s, a , ζ)}
(17)
which upon substitution in the value function VπH∗ ,ζ (s) = P ∗ H a π (a|s, ζ)Qπ ∗ ,ζ (s, a, ζ) yields the soft-min expression: ( ) X 1 H H Vπ∗ (s) = −α log exp − Qπ∗ ,ζ (s, a, ζ) . (18) α a∈A
Subsequently, the gradient expression in (14) simplifies. Corollary 1: For the objective JH (π, ζ) in (13), with optimal policy π ∗ in (17), the gradient ∇ζ JH (π ∗ , ζ)|ζ=ζ̄ = Epπ∗ ζ̄
∞ hX
γ t ∇ζ QH π ∗ ,ζ̄ (st , at , ζ)
t=0
ζ=ζ̄
i s0 ∼ µ .
(19)
Proof: Differentiating the soft-min value function (18) with respect to ζ and evaluating at ζ = ζ̄: X ∇ζ VπH∗ ,ζ (s) ζ=ζ̄ = π ∗ (a|s, ζ̄)∇ζ QH π ∗ ,ζ (s, a, ζ) ζ=ζ̄ , (20) a∈A
where ∇ζ QH π ∗ ,ζ (s, a, ζ) ζ=ζ̄ can be evaluated in a manner analogous to the ∇ζ Qπ,ζ (s, a, ζ) ζ=ζ̄ in (9). When substituted back in (20), we obtain the expression (19). B. EnvPG Actor–Critic Algorithm We now develop a model-free actor–critic algorithm that jointly optimizes the policy and environment parameter using the gradients (16) and (14). Let the parameterized policy be πθ (a|s, ζ) = P
exp{−fθ (s, a, ζ)} , exp{−fθ (s, a′ , ζ)}
(21)
a′ ∈A
where fθ (s, a, ζ) denotes the actor network and θ its weights. The actor is conditioned on the environment parameter, since the optimal policy depends on the prevailing value of ζ, and
Algorithm 1: EnvPG Actor-Critic Algorithm Input: States and actions S, A, discount factor γ, actor network fθ , critic network Qϕ , step sizes {ιk }, {βk }, {ηk }, temperature α, annealing rate υ < 1, episodes per iteration M , max outer iteration kmax , variance σ 2 , ϵ ∈ (0, 1), batch size N Initialize: θ0 , ϕ0 , ϕ− ← ϕ0 , ζ0 , k ← 0 while k ≤ kmax do Replay buffer Dk ← ∅ for episode = 1 to M do Sample initial state s0 , t ← 0 while until termination do sample at ∼ ϵ-greedy πθk (·|st , ζk ) ϵnoise ∼ N (0, σ 2 I); ζpert ← ζk + ϵnoise take action at with environment at ζpert observe st+1 and ct Dk ←Dk ∪{(st , at , st+1 , ct , ζpert )} t←t+1 sample {(s(i) , a(i) , s′(i) , c(i) , ζ (i) )}N i=1 from Dk compute TD target y (i) for all i in (22) update critic ϕk+1 via ∇ϕ L(ϕk ) (23) update target ϕ− via soft Polyak updates [1] update actor θk+1 via policy gradients in (16) if k > warmup iterations then ζk+1 ←ζk −ιk ∇ζ JH (π, ζk ) (use equation (14)) k ← k + 1; α ← υα Output: (θkmax , ζkmax )
is updated using the stochastic estimate of (16), with the score ∇θ log πθ obtained directly from (21). The critic, by contrast, must carry the two-copy structure. The gradient (14) requires capturing how QH (s, a, ζ) varies π,ζ̄ with the ζ copy — the current-step copy governing the instantaneous cost and transition probability — in a neighborhood of ζ̄. It does not require variation with respect to the ξ copy, the future-step copy of the environment parameters; there, only the value at ξ = ζ̄ is needed. This asymmetry guides the design of the critic network Qϕ (s, a, ζ), which approximates QH π,ξ (s, a, ζ): the dependence on the ζ copy is an explicit network input, while the dependence on the ξ copy is implicit, as detailed in the training procedure below. We train the critic by temporal-difference (TD) learning, building on the methodology of the deep Q-network (DQN) in [21]. At the k-th iteration, let ζk denote the baseline environment parameter. We collect data points of the form (i) (s(i) , a(i) , s′(i) , c(i) , ζpert ), where s′(i) is the successor state reached when action a(i) is taken in state s(i) , c(i) is the (i) transition cost, and ζpert = ζk + ϵnoise is the perturbed environment parameter at which the transition is observed. For each data point, the TD target is (i)
y (i) = c(i) + α log πθk (a(i) |s(i) , ζpert ) X +γ πθk (a′ |s′(i) , ζk ) Qϕ− (s′(i) , a′ , ζk ), (22) a′ ∈A
7
f1
5
User nodes {ni} 7
f3
f4
f2
3
f1
5
f5
f5
f3
f4
f2
3
1
User nodes {ni} 7
1 2
4
6
8
(a1) Alg. 1, κ = 1.0
10
User nodes {ni} 7
5
f5 f1
3
f2
f3
f4
1 2
4
6
8
(a2) Oracle, κ = 1.0
10
User nodes {ni}
5
f5 f3
3
f1
f2
f4
1 2
4
6
8
10
(b1) Algorithm 1, κ = 0.85
2
4
6
8
(b2) Oracle, κ = 0.85
10 (c) ∥∇ζ JH ∥ vs iteration k
Fig. 1. UAV placements and routing policy for Algorithm 1 and the model-based oracle for (a1)-(a2) κ = 1.0, and (b1)-(b2) κ = 0.85. User nodes are colored by their assigned UAV and arrows indicate routing under the optimal policy; (c) decaying moving average of ∥∇ζ JH ∥ as Algorithm 1 proceeds.
where Qϕ− is the target critic, a slowly-updated copy of the critic that stabilizes TD learning [21]. The first two terms (i) of y (i) depend on the perturbed parameter ζpert , while the bootstrapped term depends on the baseline ζk — realizing the (i) two-copy separation, with ζpert in the role of the current copy ζ and ζk in the role of the future copy ξ. The critic parameters are obtained by minimizing the mean-squared error, N
L(ϕ) =
2 1 X (i) Qϕ s(i) , a(i) , ζpert − y (i) . N i=1
(23)
The trained critic Qϕ (s, a, ζ) thus approximates the generalized action-value function with the future copy ξ frozen at the baseline ζk . Because the training data samples ζ only near ζk , this approximation is local: it is accurate for values of the current copy ζ within the perturbation neighborhood of ζk , which is precisely the region where the gradient ∇ζ Qϕ (s, a, ζ)|ζ=ζk is evaluated. We then use ∇ζ Qϕ (s, a, ζ)|ζ=ζk as the estimate of ∇ζ QH π,ζk (s, a, ζ)|ζ=ζk , and Qϕ (s, a, ζk ) as the estimate of QH π,ζk (s, a, ζk ), to compute ∇ζ JH (π, ζ)|ζ=ζk and ∇θ JH (πθ , ζk ). Algorithm 1 details the full procedure. V. S IMULATION We demonstrate Algorithm 1 on a synthetic UAV-assisted relay network design problem. The network consists of |N | = 271 fixed user nodes {ni } and a single destination node δ representing a base station (see Figure 1). A set of |J | = 5 UAV relay nodes {fj } with locations {yj ∈ R2 }5j=1 serve as intermediate hops between the user and destination nodes. Each message originating at a user node ni must reach δ through a sequence of hops via the UAVs, with the cost of each hop equal to the squared Euclidean distance between the nodes, approximating the communication delay. The UAV locations {yj }5j=1 , which are not fixed a priori, are stacked into the environment parameter ζ ∈ R10 . The objective is to determine ζ jointly with the routing policy so as to minimize the total expected communication cost across the network. The underlying MDP’s state space S is composed of the user node states {ni }, the UAV states {fj }, and the absorbing destination state δ, i.e., S = {{ni }, {fj }, δ}, giving |S| = 277 states in total. The action space A consists of actions that recommend message transitions to either of the five UAVs or directly to δ, i.e., A = {{fj }, δ}, giving |A| = 6 actions. The valid action set A(s) at each state is determined by the network topology: from a user node ni , the message can be
transmitted only to the UAV relays {fj }; from a UAV state fj , the message may be routed to any other UAV fj ′ ̸= fj or to the destination δ; and δ is a cost-free termination state. For the purpose of simulation we consider the following synthetic transition kernel that is parameterized by κ ∈ (0, 1]: ( κ if s′ = a , (24) pζ (s′ |s, a) = (1 − κ) qζ (s′ |s, a) otherwise where qζ (s′ |s, a) is a Gibbs distribution over the intended states valid actions A(s)\{a}: qζ (s′ |s, a) ∝ of the remaining 2 exp −η ∥xs′ − xs ∥ , with xs ∈ R2 denoting the 2D coordinates corresponding to the state s, η = 100 the concentration parameter governing the sharpness of the offtarget distribution, and κ ∈ {1.00, 0.95, 0.90, 0.85} varied across experiments. The instantaneous cost is given by cζ (s, a, s′ ) = ∥xs′ − xs ∥2 , the initial state distribution µ is uniform over the user node states, and discounting γ = 1. The above described transition kernel and cost function are used to design the emulator with which the Algorithm 1 interacts to jointly optimize the UAV placements ζ and the routing policy πθ . Note that this setup is a model-free setup as we don’t use the explicit definitions of the transition kernel and the cost function. For the purpose of simulation, the actor and critic networks are parameterized by a two-layer neural network with the state and action co-ordinates, and one-hot state and action features as inputs. The temperature α is annealed geometrically (down to ≈ 0) across outer iterations to progressively sharpen the routing policy (21). To assess the quality of the UAV placements and routing policy learned by Algorithm 1, we compare against a modelbased oracle that has full access to the transition kernel pζ (s′ |s, a) and the cost function cζ (s, a, s′ ). For a given placement ζ, the oracle can compute the exact entropyregularized value function VπH∗ ,ζ — and hence JH (π ∗ , ζ) — by soft value iteration on the recursion in (18) [3], [20], and estimates ∇ζ JH (π ∗ , ζ) by central finite differences. Gradient descent is then used to update ζ, with α annealed geometrically across iterations till the optimal policy in (17) converges to a binary value — mirroring the annealing schedule similar to the one in model-free setup. Figure 1 shows the final UAV placements obtained by Algorithm 1 and the model-based oracle for each value of κ ∈ {1.0, 0.85}; qualitatively converging to similar network geometries. In each plot, user nodes are colored according to the UAV they are routed to under the optimal policy. The arrows indicate the recommended action from each UAV
TABLE I O BJECTIVE VALUE J ∗ ACROSS κ. Method
κ = 1.00
κ = 0.95
κ = 0.90
κ = 0.85
Avg. diff.
Alg. 1 Oracle
12.99 13.00
14.00 13.83
14.80 14.76
15.60 15.27
0.9%
under the optimal policy. Note that these arrows represent the routing decision at each UAV, not necessarily the realized next hop of a communication packet, which is subject to the stochastic transition kernel. Table I provides the final objective value J ∗ , across all κ ∈ {1.0, 0.95, 0.90, 0.85}. The final objective value J ∗ achieved by both the above methods are in close agreement (across all κ values) as shown in the Table I), with an average difference ≈ 0.9% — thus, validating the proposed model-free Algorithm 1. Figure 1(c) plots the norm of the stochastic estimate of ∇ζ JH (π, ζ) against iteration for each κ; the estimate is computed over sampled minibatches, and the bold curve is a moving average over a window of 100 iterations. In all four cases the gradient norm decreases and plateaus at a small residual level — it does not vanish exactly, since it is a minibatch estimate — indicating that Algorithm 1 reaches a stationary point of the entropy-regularized objective JH (π, ζ) in (13). Since α → 0 over the course of training, JH (π, ζ) → J(π, ζ) in (1), recovering the original co-design objective. VI. C ONCLUSION AND F UTURE W ORK We established the Environment Parameter Gradient Theorem, a tractable closed-form expression for the gradient of the infinite-horizon value function with respect to continuous environment design parameters, enabled by the generalized action-value function Qπ,ξ (s, a, ζ) and its separation of the environment parameter into current-step and future-rollout copies. The resulting gradient is amenable to a modelfree setting, and makes joint policy-and-environment design possible without access to the transition kernel or cost, which we demonstrated on a UAV relay network design problem. Several directions remain open. On the theoretical side, the compatible function approximation theorem [17] — which establishes the class of critic approximators that result into an unbiased estimator for the policy gradient — has a natural counterpart here: characterizing which approximators of Qπ,ξ (s, a, ζ) leave the ζ and policy gradients unbiased. A related direction is a ζ-baseline, analogous to the advantage function [10], to reduce the variance of the gradient ∇ζ J(π, ζ) estimator. Finally, extending the framework to discrete ζ, and to constrained MDPs where ζ is subject to physical or operational constraints, would broaden its applicability to a wider class of engineering problems. R EFERENCES [1] R. S. Sutton and A. G. Barto, Reinforcement Learning: An Introduction, 2nd ed. The MIT Press, 2018. [Online]. Available: http://incompleteideas.net/book/the-book-2nd.html [2] K. S. Luck, H. B. Amor, and R. Calandra, “Data-efficient co-adaptation of morphology and behaviour with deep reinforcement learning,” in Proceedings of the Conference on Robot Learning, ser. Proceedings of Machine Learning Research, L. P. Kaelbling, D. Kragic, and K. Sugiura, Eds., vol. 100. PMLR, 30 Oct–01 Nov 2020, pp. 854–869. [Online]. Available: https://proceedings.mlr.press/v100/luck20a.html
[3] A. Srivastava and S. M. Salapaka, “Parameterized mdps and reinforcement learning problems—a maximum entropy principle-based framework,” IEEE Transactions on Cybernetics, vol. 52, no. 9, pp. 9339–9351, 2022. [4] A. Firooznia, J. Ploeg, N. Van De Wouw, and H. Zwart, “Co-design of controller and communication topology for vehicular platooning,” IEEE Transactions on Intelligent Transportation Systems, vol. 18, no. 10, pp. 2728–2739, 2017. [5] M. Garcia-Sanz, “Control co-design: an engineering game changer,” Advanced Control for Applications: Engineering and Industrial Systems, vol. 1, no. 1, p. e18, 2019. [6] L. Y. Pao, M. Pusch, and D. S. Zalkind, “Control co-design of wind turbines,” Annual Review of Control, Robotics, and Autonomous Systems, vol. 7, 2024. [7] C. Schaff, D. Yunis, A. Chakrabarti, and M. R. Walter, “Jointly learning to construct and control agents using deep reinforcement learning,” in 2019 international conference on robotics and automation (ICRA). IEEE, 2019, pp. 9798–9805. [8] M. Cauz, A. Bolland, C. Ballif, and N. Wyrsch, “Reinforcement learning for efficient design and control co-optimisation of energy systems,” in ICML 2024 AI for Science Workshop, 2024. [Online]. Available: https://openreview.net/forum?id=17tZF3ibk4 [9] D. Ha, “Reinforcement learning for improving agent design,” Artificial life, vol. 25, no. 4, pp. 352–365, 2019. [10] R. J. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning,” Machine learning, vol. 8, no. 3, pp. 229–256, 1992. [11] J. Bhatia, H. Jackson, Y. Tian, J. Xu, and W. Matusik, “Evolution gym: A large-scale benchmark for evolving soft robots,” Advances in Neural Information Processing Systems, vol. 34, pp. 2201–2214, 2021. [12] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor,” in International conference on machine learning. Pmlr, 2018, pp. 1861–1870. [13] L. Jackson, C. Walters, S. Eckersley, P. Senior, and S. Hadfield, “Orchid: optimisation of robotic control and hardware in design using reinforcement learning,” in 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2021, pp. 4911– 4917. [14] A. Bolland, I. Boukas, M. Berger, and D. Ernst, “Jointly learning environments and control policies with projected stochastic gradient ascent,” Journal of Artificial Intelligence Research, vol. 73, pp. 117– 171, 2022. [15] Z. He and M. Ciocarlie, “Morph: Design co-optimization with reinforcement learning via a differentiable hardware model proxy,” in 2024 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2024, pp. 7764–7771. [16] Y. Dai, Y. Wang, D. R. Ashley, and J. Schmidhuber, “Efficient morphology-control co-design via stackelberg proximal policy optimization,” in The Fourteenth International Conference on Learning Representations, 2026. [Online]. Available: https://openreview.net/forum?id=sJ0vOOkclw [17] R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour, “Policy gradient methods for reinforcement learning with function approximation,” Advances in neural information processing systems, vol. 12, 1999. [18] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller, “Deterministic policy gradient algorithms,” in Proceedings of the 31st International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, E. P. Xing and T. Jebara, Eds., vol. 32, no. 1. Bejing, China: PMLR, 22–24 Jun 2014, pp. 387–395. [Online]. Available: https://proceedings.mlr.press/v32/silver14.html [19] D. P. Bertsekas and J. N. Tsitsiklis, Neuro-Dynamic Programming. Belmont, MA: Athena Scientific, 1996. [20] R. Fox, A. Pakman, and N. Tishby, “Taming the noise in reinforcement learning via soft updates,” arXiv preprint arXiv:1512.08562, 2015. [21] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al., “Human-level control through deep reinforcement learning,” nature, vol. 518, no. 7540, pp. 529–533, 2015.