Null-space preference critic: dual-critic PPO for uwlab_rl, plus grad-clip and wandb-resume fixes

scripts/reinforcement_learning/rsl_rl/train.py

@@ -38,6 +38,35 @@
 parser.add_argument(
     "--ray-proc-id", "-rid", type=int, default=None, help="Automatically configured by Ray integration, otherwise None."
 )
+# --- null-space preference critic -------------------------------------------------------------
+parser.add_argument("--beta", type=float, default=None, help="Preference step budget. 0 == baseline PPO.")
+parser.add_argument(
+    "--pref_source", type=str, default=None, choices=["zero", "noise", "action_rate", "ee_height", "terms"],
+    help="Preference stream: zero (sanity A), noise (sanity B), action_rate (noise-bait probe), "
+         "terms (scripted predicates).",
+)
+parser.add_argument("--pref_noise_std", type=float, default=None, help="Std for --pref_source=noise.")
+parser.add_argument(
+    "--pref_terms", type=str, default=None,
+    help="Comma-separated RewardManager term names for --pref_source=terms.",
+)
+parser.add_argument(
+    "--critic_arch", type=str, default=None, choices=["shared", "separate"],
+    help="Second value head as a widened shared trunk, or an independent critic MLP.",
+)
+parser.add_argument("--gamma_pref", type=float, default=None, help="Discount for the preference return.")
+parser.add_argument(
+    "--pref_mask_noise", type=lambda v: v.lower() not in ("0", "false", "no"), default=None,
+    help="Layer 1: keep the preference gradient off the exploration-noise params (default true).",
+)
+parser.add_argument(
+    "--pref_detach_noise_features", action="store_true", default=False,
+    help="Layer 2: also detach the trunk features feeding the gSDE noise head in the pref surrogate.",
+)
+parser.add_argument(
+    "--projection_mode", type=str, default=None, choices=["gradient", "advantage", "sum"],
+    help="gradient = faithful null-space projection; advantage/sum = ablations.",
+)
 # append RSL-RL cli arguments
 cli_args.add_rsl_rl_args(parser)
 # append AppLauncher cli args
@@ -87,6 +116,16 @@
 
 from rsl_rl.runners import DistillationRunner, OnPolicyRunner
 
+from uwlab_rl.rsl_rl.nullspace import (
+    ActionRatePreference,
+    DualCriticOnPolicyRunner,
+    EndEffectorHeightPreference,
+    DualRewardVecEnvWrapper,
+    GaussianNoisePreference,
+    RewardManagerTermsPreference,
+    ZeroPreference,
+)
+
 from isaaclab.envs import (
     DirectMARLEnv,
     DirectMARLEnvCfg,
@@ -126,6 +165,28 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
     )
 
+    # --- null-space preference critic CLI overrides ---
+    # Applied before sanitize_rsl_rl_cfg, which only strips keys for algorithm classes it can
+    # resolve inside rsl_rl.algorithms; NullspacePPO lives in uwlab_rl, so these survive.
+    if args_cli.beta is not None:
+        agent_cfg.algorithm.beta = args_cli.beta
+    if args_cli.gamma_pref is not None:
+        agent_cfg.algorithm.gamma_pref = args_cli.gamma_pref
+    if args_cli.projection_mode is not None:
+        agent_cfg.algorithm.projection_mode = args_cli.projection_mode
+    if args_cli.pref_mask_noise is not None:
+        agent_cfg.algorithm.pref_mask_noise = args_cli.pref_mask_noise
+    if args_cli.pref_detach_noise_features:
+        agent_cfg.algorithm.pref_detach_noise_features = True
+    if args_cli.critic_arch is not None:
+        agent_cfg.policy.critic_arch = args_cli.critic_arch
+    if args_cli.pref_source is not None:
+        agent_cfg.pref_source = args_cli.pref_source
+    if args_cli.pref_noise_std is not None:
+        agent_cfg.pref_noise_std = args_cli.pref_noise_std
+    if args_cli.pref_terms is not None:
+        agent_cfg.pref_term_names = tuple(n.strip() for n in args_cli.pref_terms.split(",") if n.strip())
+
     # make config compatible with installed rsl-rl version
     agent_cfg = cli_args.sanitize_rsl_rl_cfg(agent_cfg)
 
@@ -200,13 +261,61 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         env = gym.wrappers.RecordVideo(env, **video_kwargs)
 
     # wrap around environment for rsl-rl
-    env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
+    if agent_cfg.class_name == "DualCriticOnPolicyRunner":
+        # The dual-critic path needs a second reward stream. DualRewardVecEnvWrapper reads the
+        # RewardManager's already-materialised per-term buffer rather than restructuring the
+        # manager -- OmniReset's `progress_context` term returns zeros but caches state that the
+        # reward, terminations, reset curriculum and data-collection configs all read back, so
+        # splitting or reweighting the manager silently corrupts the reward.
+        source_name = getattr(agent_cfg, "pref_source", "zero")
+        if source_name == "zero":
+            pref_source = ZeroPreference()
+        elif source_name == "noise":
+            pref_source = GaussianNoisePreference(
+                std=getattr(agent_cfg, "pref_noise_std", 1.0), seed=agent_cfg.seed
+            )
+        elif source_name == "action_rate":
+            # Noise-bait probe: a preference maximally satisfiable by shrinking exploration.
+            pref_source = ActionRatePreference()
+        elif source_name == "ee_height":
+            # High-conflict preference: fights the lift the task requires.
+            pref_source = EndEffectorHeightPreference()
+        elif source_name == "terms":
+            term_names = list(getattr(agent_cfg, "pref_term_names", ()))
+            if not term_names:
+                raise ValueError("--pref_source=terms requires --pref_terms=<comma,separated,names>")
+            pref_source = RewardManagerTermsPreference(term_names)
+        else:
+            raise ValueError(f"Unknown pref_source: {source_name}")
+        print(f"[INFO] Preference reward source: {source_name}")
+        env = DualRewardVecEnvWrapper(env, pref_source=pref_source, clip_actions=agent_cfg.clip_actions)
+    else:
+        env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
+
+    # In-job wandb with a stable run id (gen_convergence_yaml.py sets WANDB_RUN_ID + WANDB_RESUME=allow):
+    # every preemption restart resumes the SAME wandb run instead of opening a new one. On resume wandb
+    # loads the run's stored config, and rsl_rl's WandbSummaryWriter then re-sends its own config with
+    # values that change on every start (log_dir, env_cfg.log_dir, resume settings). wandb raises
+    # ConfigError on a changed value unless allow_val_change=True, which would kill the start before
+    # training -- so let config updates overwrite, in this process only.
+    if agent_cfg.logger == "wandb" and os.environ.get("WANDB_RUN_ID"):
+        import wandb.sdk.wandb_config as _wandb_config
+
+        _config_update = _wandb_config.Config.update
+
+        def _update_allow_change(self, d, allow_val_change=None):
+            return _config_update(self, d, allow_val_change=True)
+
+        _wandb_config.Config.update = _update_allow_change
+        print(f"[INFO] wandb: resuming run id {os.environ['WANDB_RUN_ID']} (config updates may overwrite)")
 
     # create runner from rsl-rl
     if agent_cfg.class_name == "OnPolicyRunner":
         runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
     elif agent_cfg.class_name == "DistillationRunner":
         runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
+    elif agent_cfg.class_name == "DualCriticOnPolicyRunner":
+        runner = DualCriticOnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
     else:
         raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
     # write git state to logs

source/uwlab_assets/uwlab_assets/__init__.py

@@ -65,7 +65,12 @@ def resolve_cloud_path(path: str) -> str:
         return path
 
     rel = _extract_relative_path(path)
-    cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "uwlab", "assets")
+    # Overridable via UWLAB_ASSET_CACHE_DIR. On a cluster ``$HOME`` is ephemeral, so the default
+    # would re-download ~7 GB of USD assets on every job; pointing this at a persistent (writable)
+    # mount makes the first job populate the cache and every later job hit it.
+    cache_dir = os.getenv("UWLAB_ASSET_CACHE_DIR") or os.path.join(
+        os.path.expanduser("~"), ".cache", "uwlab", "assets"
+    )
     local = os.path.join(cache_dir, rel)
 
     if os.path.isfile(local):

source/uwlab_rl/test/nullspace/test_grad_clip_and_pref_norm.py

@@ -0,0 +1,217 @@
+# Copyright (c) 2026, Null-space preference critic project.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the gradient-clip coupling fix (NOTES 29).
+
+The bug: rsl_rl clips the whole policy as ONE gradient vector. With a dual critic, a large
+preference value loss enters the same norm as the actor's gradient, so clipping scales the actor's
+update down -- at beta=0 too, where the preference is meant to have no effect at all. Under
+``pref_source=action_rate`` the preference value loss reached 1e5-1e6 against
+``max_grad_norm=1.0``, which throttled the actor and pinned the adaptive learning rate at its cap.
+
+``projection.py`` imports nothing but torch, so it is loaded by path (as in ``test_projection.py``)
+and none of these tests need Isaac Sim.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import io
+import pathlib
+
+import pytest
+import torch
+
+_PATH = pathlib.Path(__file__).resolve().parents[2] / "uwlab_rl" / "rsl_rl" / "nullspace" / "projection.py"
+_spec = importlib.util.spec_from_file_location("nsc_projection_gradclip", _PATH)
+_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_mod)
+partition_policy_params = _mod.partition_policy_params
+clip_grad_norm_by_group = _mod.clip_grad_norm_by_group
+
+MAX_NORM = 1.0  # the OmniReset agent config's max_grad_norm
+
+
+class _DualPolicy(torch.nn.Module):
+    """Stand-in carrying DualCriticActorCritic's real parameter names."""
+
+    def __init__(self, separate_pref_critic: bool = True) -> None:
+        super().__init__()
+        self.actor = torch.nn.Linear(4, 2)
+        self.log_std = torch.nn.Parameter(torch.zeros(2))  # gSDE noise scale -- belongs to the actor
+        self.critic = torch.nn.Linear(4, 1)
+        if separate_pref_critic:
+            self.critic_pref = torch.nn.Linear(4, 1)
+
+
+def _fill(params, value: float) -> None:
+    for p in params:
+        p.grad = torch.full_like(p, value)
+
+
+def _flat(params) -> torch.Tensor:
+    return torch.cat([p.grad.flatten() for p in params]).clone()
+
+
+def _groups(policy):
+    return partition_policy_params(policy.named_parameters())
+
+
+# -- the bug, reproduced, and the fix ------------------------------------------------------------
+
+
+def test_global_clip_lets_the_pref_critic_throttle_the_actor():
+    """Documents the bug: under one global clip the actor's update shrinks because the
+    PREFERENCE CRITIC's gradient is large -- even though the actor's own gradient is small."""
+    p = _DualPolicy()
+    g = _groups(p)
+    _fill(g["actor"] + g["critic"], 0.1)  # actor norm ~0.35: under the limit on its own
+    _fill(g["critic_pref"], 1e3)  # the action_rate preference critic's scale
+    before = _flat(g["actor"])
+    assert before.norm() < MAX_NORM
+
+    torch.nn.utils.clip_grad_norm_(p.parameters(), MAX_NORM)  # upstream behaviour
+
+    assert _flat(g["actor"]).norm() / before.norm() < 1e-2  # crushed by a different module
+
+
+def test_per_group_clip_leaves_the_actor_untouched_by_the_pref_critic():
+    p = _DualPolicy()
+    g = _groups(p)
+    _fill(g["actor"] + g["critic"], 0.1)
+    _fill(g["critic_pref"], 1e3)
+    before = _flat(g["actor"])
+
+    norms = clip_grad_norm_by_group(g, MAX_NORM)
+
+    assert torch.allclose(_flat(g["actor"]), before)  # below its own limit -> not scaled at all
+    assert _flat(g["critic_pref"]).norm() == pytest.approx(MAX_NORM, rel=1e-4)  # clipped on its own
+    assert norms["actor"] == pytest.approx(before.norm().item(), rel=1e-5)  # pre-clip norms reported
+    assert norms["critic_pref"] > 1e3
+
+
+def test_actor_update_is_invariant_to_pref_critic_scale():
+    """The property the acceptance run checks at the outcome level: at beta=0 the actor's clipped
+    gradient must not depend on how large the preference critic's gradient is. Under the global
+    clip it does; under per-group clipping it does not."""
+    per_group, global_ = [], []
+    for pref_scale in (0.0, 1.0, 1e3, 1e6):
+        for mode, out in (("per_group", per_group), ("global", global_)):
+            p = _DualPolicy()
+            torch.manual_seed(0)
+            g = _groups(p)
+            _fill(g["actor"] + g["critic"], 0.1)
+            _fill(g["critic_pref"], pref_scale)
+            if mode == "per_group":
+                clip_grad_norm_by_group(g, MAX_NORM)
+            else:
+                torch.nn.utils.clip_grad_norm_(p.parameters(), MAX_NORM)
+            out.append(_flat(g["actor"]))
+
+    assert all(torch.allclose(a, per_group[0]) for a in per_group)
+    assert not all(torch.allclose(a, global_[0]) for a in global_)
+
+
+def test_per_group_still_clips_an_oversized_actor():
+    """Per-group clipping must not silently disable clipping for the actor."""
+    p = _DualPolicy()
+    g = _groups(p)
+    _fill(g["actor"], 10.0)
+    _fill(g["critic"] + g["critic_pref"], 0.1)
+
+    clip_grad_norm_by_group(g, MAX_NORM)
+
+    assert _flat(g["actor"]).norm() == pytest.approx(MAX_NORM, rel=1e-4)
+
+
+# -- partition ------------------------------------------------------------------------------------
+
+
+def test_partition_follows_actor_parameters_and_avoids_the_prefix_trap():
+    """``critic`` is a prefix of ``critic_pref``; the preference critic must not land in the task
+    critic's group, and ``log_std`` must stay with the actor (as in ``actor_parameters``)."""
+    p = _DualPolicy()
+    ids = {k: {id(x) for x in v} for k, v in _groups(p).items()}
+
+    assert ids["actor"] == {id(p.actor.weight), id(p.actor.bias), id(p.log_std)}
+    assert ids["critic"] == {id(p.critic.weight), id(p.critic.bias)}
+    assert ids["critic_pref"] == {id(p.critic_pref.weight), id(p.critic_pref.bias)}
+
+
+def test_shared_critic_arch_leaves_pref_group_empty_and_is_safe():
+    """``critic_arch='shared'`` has no separate preference critic; params without a gradient are
+    skipped rather than erroring."""
+    p = _DualPolicy(separate_pref_critic=False)
+    g = _groups(p)
+    assert g["critic_pref"] == []
+    _fill(g["critic"], 0.1)  # actor deliberately left with grad=None
+
+    norms = clip_grad_norm_by_group(g, MAX_NORM)
+
+    assert norms["critic_pref"] == 0.0 and norms["actor"] == 0.0 and norms["critic"] > 0.0
+
+
+# -- preference reward normalisation -------------------------------------------------------------
+
+_networks = pytest.importorskip("rsl_rl.networks")
+EDVN = _networks.EmpiricalDiscountedVariationNormalization
+
+
+def test_pref_reward_normalizer_brings_action_rate_scale_rewards_to_order_one():
+    torch.manual_seed(0)
+    n = EDVN(shape=1, gamma=0.99)
+    n.train()
+    for _ in range(300):
+        raw = -(torch.rand(4096, 1) * 3e3)  # -||a - a_prev||^2 at the magnitudes that were observed
+        out = n(raw)
+
+    assert raw.abs().mean() > 1e3
+    assert 1e-3 < out.abs().mean().item() < 10.0
+
+
+def test_zero_preference_stays_exactly_zero():
+    """``pref_source='zero'`` must be untouched: no NaN from a vanishing std, no drift off zero."""
+    n = EDVN(shape=1, gamma=0.99)
+    n.train()
+    for _ in range(200):
+        out = n(torch.zeros(4096, 1))
+
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, torch.zeros_like(out))
+
+
+def test_normalizer_works_under_inference_mode_and_survives_a_checkpoint_round_trip():
+    """rsl_rl collects rollouts inside ``torch.inference_mode()`` (on_policy_runner.py:101), so the
+    normaliser's buffers are rewritten there. They must still be readable for logging, serialisable
+    by ``torch.save``, and loadable into a freshly built policy on resume."""
+    n = EDVN(shape=1, gamma=0.99)
+    n.train()
+    with torch.inference_mode():
+        for _ in range(20):
+            n(-(torch.rand(512, 1) * 1e3))
+    scale = float(n.emp_norm._std)  # read outside inference mode, as update() logging does
+    assert scale > 1.0
+
+    buf = io.BytesIO()
+    torch.save(n.state_dict(), buf)  # what the runner's save() does
+    buf.seek(0)
+    fresh = EDVN(shape=1, gamma=0.99)
+    fresh.load_state_dict(torch.load(buf))  # what a resumed run does, outside inference mode
+    assert float(fresh.emp_norm._std) == pytest.approx(scale)
+
+    fresh.train()
+    with torch.inference_mode():
+        out = fresh(-(torch.rand(512, 1) * 1e3))
+    assert torch.isfinite(out).all()
+
+
+def test_normalizer_on_the_policy_is_checkpointed_and_adds_no_parameters():
+    """Attached to the policy it rides along in ``policy.state_dict()`` -- and, holding only
+    buffers, it cannot enter a clip group or the optimizer."""
+    p = _DualPolicy()
+    p.pref_reward_normalizer = EDVN(shape=1, gamma=0.99)
+
+    assert any(k.startswith("pref_reward_normalizer.emp_norm.") for k in p.state_dict())
+    assert not any(n.startswith("pref_reward_normalizer") for n, _ in p.named_parameters())
+    grouped = {id(x) for v in _groups(p).values() for x in v}
+    assert grouped == {id(x) for x in p.parameters()}  # every parameter grouped exactly once