# Decoupled Environment Gradient for Joint Policy and Simulator Learning

- ID: 819
- Canonical URL: https://synthcore.org/idea/819/decoupled-environment-gradient-for-joint-policy-and-simulato
- API JSON: https://synthcore.org/api/idea/819.json
- API Markdown: https://synthcore.org/api/idea/819.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.12590](https://arxiv.org/abs/2607.12590)
- Category: optimization
- Solves: sample-efficiency, accuracy, stability
- ML areas: rl, world-model, optimizer, training-dynamics
- Math tags: control-theory, stochastic-processes, optimization, dynamical-systems
- Ratings: usefulness 7/10; difficulty 5/10; novelty 7/10

## Idea description

Augment a neural policy with differentiable environment or data-generation parameters and optimize both using the environment-parameter policy-gradient theorem. The current transition is differentiated with respect to the design parameter, while the continuation value is evaluated under a frozen copy of that parameter; this isolates the local causal effect and avoids repeatedly differentiating through arbitrarily long rollouts. Suitable applications include learnable domain randomization, differentiable augmentation schedules, simulator calibration, and world-model co-design.

## Mathematical statement

Let an MDP have state s, action a, next state s', environment parameter zeta, transition kernel p_zeta(s'|s,a), cost c_zeta(s,a,s'), policy pi_theta(a|s), and discount gamma in (0,1). Define V_{theta,xi}(s)=E_{pi_theta,p_xi}[sum_{t>=0} gamma^t c_xi(s_t,a_t,s_{t+1}) | s_0=s], where xi is the parameter used in future rollouts. Define Q_{theta,xi}(s,a,zeta)=sum_{s'} p_zeta(s'|s,a)[c_zeta(s,a,s')+gamma V_{theta,xi}(s')]. The environment-gradient theorem states that the derivative of J(theta,zeta)=E_{s_0~rho_0}[V_{theta,zeta}(s_0)] is the discounted state-occupancy expectation of the local derivative of Q: gradient_zeta J = E_{s~d_{theta,zeta},a~pi_theta}[gradient_zeta Q_{theta,xi}(s,a,zeta)] with xi set equal to zeta only after differentiation. Expanding the derivative gives a transition-score term weighted by immediate cost plus continuation value, and a direct derivative of the cost. Here d_{theta,zeta}(s)=sum_{t>=0} gamma^t Pr(s_t=s) is the unnormalized discounted occupancy. The essential property is that gradient_zeta does not differentiate through V_{theta,xi}; xi is held fixed, preventing duplicated or unstable long-horizon derivatives.

## Key formulas

- $$Q_{\theta,\xi}(s,a,\zeta)=\sum_{s'}p_{\zeta}(s'\mid s,a)\left[c_{\zeta}(s,a,s')+\gamma V_{\theta,\xi}(s')\right],\qquad V_{\theta,\xi}(s)=\mathbb{E}_{\pi_\theta,p_\xi}\left[\sum_{k=0}^{\infty}\gamma^k c_\xi(s_k,a_k,s_{k+1})\mid s_0=s\right].$$
- $$\nabla_\zeta J(\theta,\zeta)=\mathbb{E}_{s\sim d_{\theta,\zeta},\,a\sim\pi_\theta}\left[\nabla_\zeta Q_{\theta,\xi}(s,a,\zeta)\big|_{\xi=\zeta}\right].$$
- $$\nabla_\zeta Q_{\theta,\xi}(s,a,\zeta)=\sum_{s'}\left[\nabla_\zeta p_\zeta(s'\mid s,a)\left(c_\zeta(s,a,s')+\gamma V_{\theta,\xi}(s')\right)+p_\zeta(s'\mid s,a)\nabla_\zeta c_\zeta(s,a,s')\right].$$
- $$\widehat{g}_\zeta=\frac{1}{B}\sum_{i=1}^{B}\left[\nabla_\zeta c_\zeta(s_i,a_i,s'_i)+\nabla_\zeta\log p_\zeta(s'_i\mid s_i,a_i)\,\widehat{Q}_{\theta,\xi}(s_i,a_i,\zeta)\right],\qquad \zeta\leftarrow\Pi_{\mathcal Z}(\zeta-\alpha_\zeta\widehat{g}_\zeta).$$

## Implementation notes

1. Integration point: implement a neural policy pi_theta(a|s) together with a differentiable environment generator p_zeta and cost c_zeta. Start with a small differentiable simulator or learnable augmentation environment, such as a two-dimensional navigation task whose obstacle locations, friction, or observation-noise scale are controlled by zeta. Use zeta as a constrained nn.Parameter and project it into a valid interval after every update. 2. Pseudocode: collect B transitions (s_i,a_i,s'_i,c_i) using the current theta and zeta; estimate a critic Q_hat_theta_xi with target rollouts or temporal-difference learning using a detached copy xi=zeta.detach(); compute the local cost derivative grad_zeta c_zeta and transition score grad_zeta log p_zeta(s'_i|s_i,a_i); form g_zeta as the average of grad_zeta c plus the score multiplied by Q_hat; update theta with the ordinary policy-gradient loss and update zeta with zeta <- projection(zeta-alpha_zeta*g_zeta). Every K steps copy xi <- zeta.detach(). For reparameterizable simulators, replace the score term with the pathwise derivative of the one-step transition while still stopping gradients through future rollout parameters. 3. Computed versus estimated: derivatives of p_zeta and c_zeta are obtained by autodiff when the simulator is differentiable; discounted occupancy, Q, and the expectation are estimated from replay trajectories. If p_zeta is unknown, fit a probabilistic one-step model and use its analytic log-density gradient. 4. First cheap experiment: use a tabular or two-layer MLP policy on a four-state navigation MDP, with zeta controlling slip probability and collision cost. Compare policy-only learning, naive backpropagation through full unrolled episodes, and the decoupled estimator. Sweep horizon H and finite-difference zeta perturbations to measure bias. The quantitative prediction is that the decoupled gradient agrees with finite differences within 10% at short horizons and its error remains approximately constant as H grows, whereas full-rollout differentiation error or variance grows with H. Joint optimization should reduce final cost below policy-only training when the optimum requires changing slip or collision parameters. Also test the predicted stationary condition g_zeta approximately zero and verify that the learned zeta is within 20% of a grid-search optimum in the small MDP.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: Built an exact two-state differentiable MDP verifier implementing the decoupled local environment-gradient theorem with a frozen continuation value. The mechanism manifested: matched-copy gradients agreed with finite differences to about 1e-11–2e-9, truncated-tail error decayed geometrically (fitted ratio 0.766, below the predicted gamma=0.92 bound), and stale-copy bias was approximately linear in the copy offset (R²=0.9968) while vanishing at zero offset. Environment optimization reduced cost from 8.5405 to 7.56210, matched the grid optimum z=-0.357 within 0.00030, and reached gradient 7.25e-8.

### Mechanism check

- Verdict: Built an exact two-state differentiable MDP verifier implementing the decoupled local environment-gradient theorem with a frozen continuation value. The mechanism manifested: matched-copy gradients agreed with finite differences to about 1e-11–2e-9, truncated-tail error decayed geometrically (fitted ratio 0.766, below the predicted gamma=0.92 bound), and stale-copy bias was approximately linear in the copy offset (R²=0.9968) while vanishing at zero offset. Environment optimization reduced cost from 8.5405 to 7.56210, matched the grid optimum z=-0.357 within 0.00030, and reached gradient 7.25e-8.
- Confidence: 9/10
- Limitations: The experiment uses an analytically solved two-state MDP and a fixed tabular policy, not a neural policy, sampled critic, replay trajectories, or a genuinely joint policy/environment learning loop. It does not compare variance or wall-clock performance against full-rollout autodiff, and the transition spectral radius is exactly one, so the observed tail ratio is governed by discounting rather than a strict subunit transition contraction.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 1
- Paired wins: 0/8
- Benchmark verdict: no measurable effect

## Artifacts

- [bench_experiment.py](https://synthcore.org/code/333/bench_experiment.py)
- [bench_report.json](https://synthcore.org/code/333/bench_report.json)
- [gradient_experiment.py](https://synthcore.org/code/333/gradient_experiment.py)
- [optimization_results.json](https://synthcore.org/code/333/optimization_results.json)
- [report.md](https://synthcore.org/code/333/report.md)
- [report_bench_2026-08-31T012832.md](https://synthcore.org/code/333/report_bench_2026-08-31T012832.md)
- [results.json](https://synthcore.org/code/333/results.json)
- [Download all files as ZIP](https://synthcore.org/download/333)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.