Null-space preference critic: dual-critic PPO for uwlab_rl, plus grad-clip and wandb-resume fixes - #43
Draft
daphnechen wants to merge 7 commits into
Draft
Null-space preference critic: dual-critic PPO for uwlab_rl, plus grad-clip and wandb-resume fixes#43daphnechen wants to merge 7 commits into
daphnechen wants to merge 7 commits into
Conversation
Adds uwlab_rl/rsl_rl/nullspace/:
projection.py pure-torch null-space projection + diagnostics (unit-tested, no sim needed)
dual_actor_critic second value head; shared-trunk or separate-critic arch
dual_storage parallel reward/value/return/advantage buffers, two GAE passes,
each stream normalised over its own full batch
nullspace_ppo one importance ratio and one clip shared by both surrogates; the
preference gradient is projected orthogonally to the task gradient
before it reaches the actor; per-stream time-out bootstrapping
reward_split task/preference split read off RewardManager._step_reward, with a
pluggable preference source (zero / gaussian noise / named terms)
runner resolves the above by name (upstream _construct_algorithm uses eval()
in its own module namespace)
Why the reward split reads _step_reward instead of restructuring the RewardManager:
omnireset's `progress_context` term returns zeros but caches the goal distances that r_dist,
r_success, three termination terms, the reset-distribution monitor and the data-collection
configs all read back off it -- and RewardManager skips any term whose weight is 0.0 without
calling it. Splitting, reordering or reweighting the manager corrupts the reward silently.
Also ports the parent checkout's uncommitted UWLAB_TMP_DIR fix in get_temp_dir(); without it
a fresh clone dies on /tmp/uwlab permissions after ~2 min of Isaac boot on shared machines.
Registers OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-Nullspace-v0 (same env as the
baseline task, so a beta=0 run is comparable line-for-line) and threads the null-space
options through train.py.
$HOME is ephemeral on Singularity, so the hard-coded ~/.cache/uwlab/assets would re-download ~7 GB of USD/HDR assets on every job. Pointing this at a persistent writable mount lets the first job populate any cache miss and every later job hit it.
…leak fix) Every scripted manner preference (mechanical power, EE speed, action smoothness) is monotonically improved by shrinking action noise, while perturbing noise near a local optimum of the mean policy is ~second order in task return. Large first-order preference gradient against ~zero first-order task gradient means the projector does not merely permit exploration collapse -- it selects for it. That is entropy collapse arriving through the exact channel built to find task-neutral directions. gSDE noise is an optimisation parameter, not a deployed behavioural property (the policy is deployed on the mean action), so masking gives up no legitimate preference. Layer 1 (pref_mask_noise, default ON): exclude the noise params from g_pref and run the projection *restricted to* the mean-action subspace. Restriction rather than zero-padding matters: the projector subtracts a multiple of the full g_task, which has components on the excluded coords, so zero-padding feeds a correction back onto exactly the parameters being protected. Covered by test_masking_is_not_equivalent_to_zeroing_gpref_in_full_space. Layer 2 (pref_detach_noise_features, opt-in): Layer 1 does NOT close the leak on its own -- gSDE's variance is mm(actor[:-1](obs)**2, exp(log_std)**2), so the preference can still shrink exploration by reshaping the trunk features feeding the noise head. Detaching both inputs to the variance inside the preference surrogate closes it. The resulting ratio is value-identical to the task ratio (detach changes no values), so "one ratio, one clip" still holds; a runtime assertion enforces that. Also: - ActionRatePreference: the noise-bait probe. r_pref = -||a - a_prev||^2, maximally satisfiable by shrinking noise and barely satisfiable otherwise. Sanity run B (Gaussian) is blind to this failure -- an unsystematic gradient does not preferentially shrink noise -- so the Gaussian run is retained but is no longer the sharp test. - Guard metrics in EVERY arm including beta=0 (guard_entropy, guard_noise_std, and the active scoping flags), since the collapse is only visible as drift relative to the beta=0 reference. - Entropy bonus documented as staying on the task side: it is a regulariser on the optimisation, not a preference about behaviour, and must keep its unrestricted path to the noise params.
- projection_mode now has three distinct behaviours (`sum` and `advantage` previously shared
one code path). `sum` is arm B': dual critic, separately normalised advantages, NO
projection -- the arm that separates scale-invariance from the projection itself. Without it
a C-beats-B result is consistent with either being the contribution.
- pref_removed_frac / cos_before logged as a TIME SERIES in every arm, including beta=0 and the
unprojected ablations, via a diagnostic-only backward that applies nothing. Alignment moves
over training and the late-training regime (g_task -> 0) is where the self-scheduling claim
lives; a converged scalar is not evidence for it.
- EndEffectorHeightPreference: a deliberately high-conflict preference ("keep the EE low" fights
the lift the task requires). The action-rate bait measures pref_removed_frac ~ 1e-3, i.e. the
constraint is barely binding -- in high-dimensional parameter space near-orthogonality is the
default, so on such preferences the projected and unprojected arms are near-identical BY
CONSTRUCTION and Phase 2 would produce no result however well the code works. Screen candidate
preferences on pref_removed_frac and span a range.
guard_noise_std is a product -- realised std ~ ||f(obs)|| * sigma -- and Layer 1 masks only sigma. The aggregate therefore cannot distinguish 'the trunk reshaped its features' from 'sigma shrank', which are different findings with different remedies. Log both factors (guard_feat_norm, guard_sigma) per iteration in every arm so the Layer 2 decision is mechanical rather than a judgement call. Also asserts Layer 1's claim at runtime: the combined gradient must equal g_task EXACTLY on the excluded noise coordinates. If sigma moves in a masked arm that is an implementation bug, not a leak, and the two must never be confused when reading the guard metrics.
…he preference reward
The beta=0 null was not a null. NullspacePPO inherited upstream rsl_rl's single global
clip_grad_norm_ over the whole policy, so the preference critic's gradient entered the same
norm as the actor's. The preference critic regresses onto unnormalised returns, and with
pref_source=action_rate its value loss reached 1e5-1e6 against max_grad_norm=1.0. That coupled
the preference critic to the actor even at beta=0 -- and because -||a - a_prev||^2 scales with
action noise, the coupling was strongest in exactly the arms that preserve exploration
(NOTES 29).
Two changes, both on by default:
- grad_clip_mode="per_group": clip the actor, the task critic and the preference critic each
on its own budget. "global" keeps the old behaviour so pre-fix runs can be reproduced.
- normalize_pref_reward=True: scale the preference reward by a running std of its discounted
sum (rsl_rl's EmpiricalDiscountedVariationNormalization) before it reaches the preference
critic. Per-stream advantage normalisation already makes the actor's preference signal
scale-free, so this fixes the critic's scale without changing what the actor sees. The task
stream is not normalised. The normaliser is registered on the policy, so it is checkpointed
and survives auto-resume.
Pre-clip gradient norms are now logged per group (grad_norm/{actor,critic,critic_pref}, or
grad_norm/global), so a cross-group throttle is directly visible rather than inferred.
Also corrects two docstrings that promised what the bug broke: critic_arch="separate" did not
by itself guarantee the beta=0 overlay, and beta=0 reproduced baseline PPO only for
pref_source="zero".
Validated:
- 10 new unit tests (test_grad_clip_and_pref_norm.py): the bug under a global clip; actor
isolation, and invariance of the actor gradient to preference-critic scale, under per-group
clipping; the critic_pref/critic prefix trap; normaliser behaviour on large and on zero
rewards; updates under torch.inference_mode; a torch.save/load round trip.
- Local A/B smoke, beta=0 with pref_source=action_rate, 6 iterations at 256 envs. Legacy global
clip: preference value loss 7.1e3 rising to 1.9e4, global grad norm 126-478 against
max_grad_norm=1.0. Fix: preference value loss 0.17-6.5, actor grad norm 0.27-0.53 (never
clipped), normaliser buffers present in the checkpoint.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BXF4A9BzFr7PX66G1x72yu
Amulet specs set a stable WANDB_RUN_ID with WANDB_RESUME=allow so a job that is preempted and restarted continues one wandb run instead of opening a new one per start. On resume wandb loads the run's stored config, and rsl_rl's WandbSummaryWriter then re-sends its own config containing values that differ on every start (log_dir, env_cfg.log_dir, resume settings). wandb raises ConfigError on a changed value unless allow_val_change=True, which killed the restart before training began. Patch Config.update in this process only, and only when resuming an explicitly named run, so those config updates overwrite instead of raising. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BXF4A9BzFr7PX66G1x72yu
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a dual-critic PPO variant to
uwlab_rlin which a preference reward gets its own critic andits gradient is projected into the null space of the task gradient:
The intent is that preference pressure selects among success-optimal behaviours rather than trading
against task success. Two fixes found while running it are included.
No new dependencies. No linked issue — this was developed as a research branch rather than from a
filed proposal, and I'm opening it as a draft for that reason.
Empirical status, stated up front
The projection's central claim did not survive its control, and a reviewer should know that
before reading the code. Measured on peg-in-hole (OmniReset, UR5e + Robotiq 2F-85, 4,000 iterations,
8 seeds per dose, 16 at the operating point), the projection was compared against naive mixing —
adding
beta * g_prefdirectly, no projection, available in this branch asprojection_mode="sum"— at matched compliance rather than matched beta, since the two modesreach different compliance at the same beta:
Both modes are safe at roughly +0.13 compliance and lose the task by roughly +0.175; the cliff sits
in the same place with the projection and without it. On this evidence the dose, not the
projection, selects the safe operating point.
To be precise about what that does and does not mean: this is no evidence of an advantage, and it
is not evidence of equivalence — two doses that both lose nothing cannot discriminate between
them, and the single pairing that favours the projection is p = 0.20 and not exactly
compliance-matched.
So the honest framing of this contribution is scale-free preference control with a null-space
option, not a demonstrated improvement from projecting.
What's here
source/uwlab_rl/.../rsl_rl/nullspace/—nullspace_ppo.py,dual_actor_critic.py,dual_storage.py,projection.py,reward_split.py,runner.pysource/uwlab_rl/uwlab_rl/rsl_rl/nullspace_cfg.py— config surface, includingprojection_modein("gradient", "advantage", "sum");"sum"is the no-projection baselinesource/uwlab_tasks/.../omnireset/config/ur5e_robotiq_2f85/— agent config for the abovesource/uwlab_rl/test/nullspace/test_grad_clip_and_pref_norm.py— tests for the clipping andpreference-normalisation behaviour
scripts/.../rsl_rl/train.py— resume the same wandb run across preemption restarts whenWANDB_RUN_IDis set, instead of opening a new run per restartsource/uwlab_assets/uwlab_assets/__init__.py— make the asset cache dir overridable viaUWLAB_ASSET_CACHE_DIRA note on size
This is +1,897 lines across 14 files, which is larger than the guidance in the PR template asks
for. Two commits are independently useful and trivially splittable if you'd prefer them separately:
62eb368—UWLAB_ASSET_CACHE_DIRoverride (3 lines, unrelated to the rest)64768df— wandb resume across preemption restarts (17 lines, unrelated to the rest)Happy to split those out, or to drop the research code entirely and land only the two fixes.
The branch is 7 commits ahead of
mainand 4 behind, but no file is touched on both sides of themerge-base, so it merges cleanly without a rebase.
Type of change
Checklist
pre-commitchecks with./uwlab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists thereUnticked items are genuinely not done rather than overlooked: the
pre-committoolchain isn'tinstalled in this environment, so I've left the format check to CI; and the changelog/version bumps
and
CONTRIBUTORS.mdentry are held until you say whether you want this branch at all in itscurrent shape. All new files do carry the standard copyright + SPDX header, and all changed Python
files compile.
🤖 Generated with Claude Code
https://claude.ai/code/session_01BXF4A9BzFr7PX66G1x72yu