scripts/train_dropbear.py

#!/usr/bin/env python3
"""Register Dropbear tasks, then run Isaac Lab's stock RSL-RL trainer.

The optional ``--warm-start-std`` flag preserves the actor, critic, and
normalizers from a checkpoint while resetting the optimizer and the policy's
collapsed exploration standard deviation.  The flag is handled here and
removed before Isaac Lab parses its own arguments.
"""

from __future__ import annotations

import os
import runpy
import sys
import tempfile
from pathlib import Path
from typing import Any


WORKSPACE_ROOT = Path(__file__).resolve().parents[1]
TRAIN_SCRIPT = WORKSPACE_ROOT / "IsaacLab" / "scripts" / "reinforcement_learning" / "rsl_rl" / "train.py"


def _pop_float_arg(name: str) -> float | None:
    """Remove a wrapper-only ``--name value`` argument and return its value."""
    if name not in sys.argv:
        return None
    index = sys.argv.index(name)
    try:
        value = float(sys.argv[index + 1])
    except (IndexError, ValueError) as exc:
        raise ValueError(f"{name} requires a numeric value") from exc
    del sys.argv[index : index + 2]
    return value


def _pop_int_arg(name: str) -> int | None:
    """Remove a wrapper-only ``--name value`` argument and return its value."""
    if name not in sys.argv:
        return None
    index = sys.argv.index(name)
    try:
        value = int(sys.argv[index + 1])
    except (IndexError, ValueError) as exc:
        raise ValueError(f"{name} requires an integer") from exc
    del sys.argv[index : index + 2]
    return value


def _pop_string_arg(name: str) -> str | None:
    """Remove a wrapper-only ``--name value`` argument and return its value."""
    if name not in sys.argv:
        return None
    index = sys.argv.index(name)
    try:
        value = sys.argv[index + 1]
    except IndexError as exc:
        raise ValueError(f"{name} requires a value") from exc
    del sys.argv[index : index + 2]
    return value


def _pop_float_pair_arg(name: str) -> tuple[float, float] | None:
    """Remove a wrapper-only ``--name low high`` argument."""
    if name not in sys.argv:
        return None
    index = sys.argv.index(name)
    try:
        low = float(sys.argv[index + 1])
        high = float(sys.argv[index + 2])
    except (IndexError, ValueError) as exc:
        raise ValueError(f"{name} requires two numeric values") from exc
    del sys.argv[index : index + 3]
    return low, high


def _enable_exploration_warm_start(std: float) -> None:
    """Patch checkpoint loading to retain weights but restart exploration."""
    if not 0.01 <= std <= 2.0:
        raise ValueError("--warm-start-std must be between 0.01 and 2.0")

    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load

    def load_weights_without_optimizer(self, path, load_cfg=None, strict=True, map_location=None):
        selected = {
            "actor": True,
            "critic": True,
            "optimizer": False,
            "iteration": True,
            "rnd": False,
        }
        infos = original_load(
            self,
            path,
            load_cfg=selected,
            strict=strict,
            map_location=map_location,
        )
        policy = self.alg.get_policy()
        with torch.no_grad():
            policy.distribution.std_param.fill_(std)
        print(
            f"[WARM START] Preserved actor/critic at iteration "
            f"{self.current_learning_iteration}; reset optimizer and action std to {std:g}."
        )
        return infos

    OnPolicyRunner.load = load_weights_without_optimizer


def _enable_actor_base_lin_vel_checkpoint_adapter() -> None:
    """Expand a 53-D actor checkpoint for the optional 56-D feedback input.

    The new base-linear-velocity channels are prepended to the policy
    observation. Their normalizer starts at mean=0/std=1 and their first-layer
    weights start at exactly zero, so conversion preserves the source actor's
    output before learning resumes. The optimizer must be reset because its
    first-layer moment tensors have the old shape.
    """
    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load

    def load_with_feedback_adapter(self, path, load_cfg=None, strict=True, map_location=None):
        checkpoint = torch.load(path, weights_only=False, map_location="cpu")
        actor_state = checkpoint["actor_state_dict"]
        first_layer = actor_state["mlp.0.weight"]
        old_width = int(first_layer.shape[1])
        expected_width = int(self.alg.get_policy().mlp[0].in_features)

        if old_width == expected_width:
            return original_load(
                self,
                path,
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        if old_width + 3 != expected_width:
            raise ValueError(
                "Actor feedback adapter expected a three-channel expansion, "
                f"but checkpoint/model widths are {old_width}/{expected_width}."
            )

        zeros = torch.zeros(
            first_layer.shape[0],
            3,
            dtype=first_layer.dtype,
            device=first_layer.device,
        )
        actor_state["mlp.0.weight"] = torch.cat((zeros, first_layer), dim=1)
        for key in ("obs_normalizer._mean", "obs_normalizer._var", "obs_normalizer._std"):
            tensor = actor_state[key]
            fill = 0.0 if key.endswith("_mean") else 1.0
            prefix = torch.full(
                (tensor.shape[0], 3),
                fill,
                dtype=tensor.dtype,
                device=tensor.device,
            )
            actor_state[key] = torch.cat((prefix, tensor), dim=1)

        temporary_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as handle:
                temporary_path = Path(handle.name)
            torch.save(checkpoint, temporary_path)
            print(
                "[OBS ADAPTER] Expanded actor observations "
                f"{old_width}->{expected_width}; prepended zero-weight "
                "base linear velocity feedback."
            )
            return original_load(
                self,
                str(temporary_path),
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        finally:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)

    OnPolicyRunner.load = load_with_feedback_adapter


def _enable_pose_reference_checkpoint_adapter() -> None:
    """Append two independent pose-reference channels to actor and critic.

    Both first-layer columns start at zero and both normalizers start at
    mean=0/std=1. Consequently the adapted 58-D actor is exactly equivalent
    to its source 56-D actor until PPO learns to use ``[depth, direction]``.
    The critic is expanded from 70 to 72 inputs by the same construction.
    """
    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load

    def _append_channels(state, expected_width: int, model_name: str) -> bool:
        first_layer = state["mlp.0.weight"]
        old_width = int(first_layer.shape[1])
        if old_width == expected_width:
            return False
        if old_width + 2 != expected_width:
            raise ValueError(
                f"{model_name} pose adapter expected a two-channel expansion, "
                f"but checkpoint/model widths are {old_width}/{expected_width}."
            )

        zeros = torch.zeros(
            first_layer.shape[0],
            2,
            dtype=first_layer.dtype,
            device=first_layer.device,
        )
        state["mlp.0.weight"] = torch.cat((first_layer, zeros), dim=1)
        for key in ("obs_normalizer._mean", "obs_normalizer._var", "obs_normalizer._std"):
            tensor = state[key]
            fill = 0.0 if key.endswith("_mean") else 1.0
            suffix = torch.full(
                (tensor.shape[0], 2),
                fill,
                dtype=tensor.dtype,
                device=tensor.device,
            )
            state[key] = torch.cat((tensor, suffix), dim=1)
        return True

    def load_with_pose_adapter(self, path, load_cfg=None, strict=True, map_location=None):
        checkpoint = torch.load(path, weights_only=False, map_location="cpu")
        actor_changed = _append_channels(
            checkpoint["actor_state_dict"],
            int(self.alg.actor.mlp[0].in_features),
            "Actor",
        )
        critic_changed = _append_channels(
            checkpoint["critic_state_dict"],
            int(self.alg.critic.mlp[0].in_features),
            "Critic",
        )
        if not actor_changed and not critic_changed:
            return original_load(
                self,
                path,
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        if actor_changed != critic_changed:
            raise ValueError("Pose adapter requires actor and critic to expand together.")

        temporary_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as handle:
                temporary_path = Path(handle.name)
            torch.save(checkpoint, temporary_path)
            print(
                "[POSE ADAPTER] Expanded actor 56->58 and critic 70->72; "
                "appended zero-weight [depth, direction] reference channels."
            )
            return original_load(
                self,
                str(temporary_path),
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        finally:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)

    OnPolicyRunner.load = load_with_pose_adapter


def _enable_appended_observation_checkpoint_adapter(label: str) -> None:
    """Append zero-initialized observation channels to actor and critic.

    Terrain scans are appended after the existing proprioceptive/reference
    observations.  Zero first-layer columns preserve the source policy exactly
    while PPO learns how to use the newly normalized exteroceptive inputs.
    """
    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load

    def _append(state, expected_width: int, model_name: str) -> int:
        first_layer = state["mlp.0.weight"]
        old_width = int(first_layer.shape[1])
        added = expected_width - old_width
        if added == 0:
            return 0
        if added < 0:
            raise ValueError(
                f"{model_name} {label} adapter cannot shrink observations "
                f"{old_width}->{expected_width}."
            )
        zeros = torch.zeros(
            first_layer.shape[0],
            added,
            dtype=first_layer.dtype,
            device=first_layer.device,
        )
        state["mlp.0.weight"] = torch.cat((first_layer, zeros), dim=1)
        for key in (
            "obs_normalizer._mean",
            "obs_normalizer._var",
            "obs_normalizer._std",
        ):
            tensor = state[key]
            fill = 0.0 if key.endswith("_mean") else 1.0
            suffix = torch.full(
                (tensor.shape[0], added),
                fill,
                dtype=tensor.dtype,
                device=tensor.device,
            )
            state[key] = torch.cat((tensor, suffix), dim=1)
        return added

    def load_with_appended_observations(
        self,
        path,
        load_cfg=None,
        strict=True,
        map_location=None,
    ):
        checkpoint = torch.load(path, weights_only=False, map_location="cpu")
        actor_added = _append(
            checkpoint["actor_state_dict"],
            int(self.alg.actor.mlp[0].in_features),
            "Actor",
        )
        critic_added = _append(
            checkpoint["critic_state_dict"],
            int(self.alg.critic.mlp[0].in_features),
            "Critic",
        )
        if actor_added == 0 and critic_added == 0:
            return original_load(
                self,
                path,
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        if actor_added != critic_added:
            raise ValueError(
                f"{label} adapter requires equal actor/critic expansion, got "
                f"{actor_added}/{critic_added}."
            )

        temporary_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(
                suffix=".pt",
                delete=False,
            ) as handle:
                temporary_path = Path(handle.name)
            torch.save(checkpoint, temporary_path)
            print(
                f"[OBS ADAPTER] Appended {actor_added} zero-weight {label} "
                "channels to actor and critic."
            )
            return original_load(
                self,
                str(temporary_path),
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        finally:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)

    OnPolicyRunner.load = load_with_appended_observations


def _enable_pose_reference_normalizer_calibration(
    reference_mean: tuple[float, float] = (0.375, 0.0),
    reference_var: tuple[float, float] = (0.18616071428571428, 0.375),
    label: str = "POSE",
) -> None:
    """Calibrate the two reference channels without disturbing legacy inputs.

    The adapted policy inherits a single observation-normalizer count exceeding
    one billion samples. Appending pose channels with mean=0/std=1 therefore
    prevents their statistics from adapting on a useful timescale. Replace
    only the final ``[depth, direction]`` statistics with their analytical
    values for the 16-second pose cycle while preserving all legacy channels.
    """
    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load
    def _calibrate(state, model_name: str) -> None:
        mean = state["obs_normalizer._mean"]
        var = state["obs_normalizer._var"]
        std = state["obs_normalizer._std"]
        if mean.shape[-1] < 2:
            raise ValueError(f"{model_name} has no pose-reference channels")
        new_mean = torch.tensor(
            reference_mean, dtype=mean.dtype, device=mean.device
        )
        new_var = torch.tensor(
            reference_var, dtype=var.dtype, device=var.device
        )
        new_std = torch.sqrt(new_var)
        old_mean = mean[..., -2:].flatten().clone()
        old_std = std[..., -2:].flatten().clone()
        pose_weights = state["mlp.0.weight"][:, -2:].clone()

        # Preserve the policy exactly across the affine normalization change:
        # W_old * ((x-m_old)/s_old) + b_old
        # == W_new * ((x-m_new)/s_new) + b_new.
        state["mlp.0.weight"][:, -2:] = pose_weights * (
            new_std / old_std
        ).unsqueeze(0)
        state["mlp.0.bias"] += pose_weights @ (
            (new_mean - old_mean) / old_std
        )
        mean[..., -2:] = new_mean
        var[..., -2:] = new_var
        std[..., -2:] = new_std

    def load_with_pose_calibration(
        self, path, load_cfg=None, strict=True, map_location=None
    ):
        checkpoint = torch.load(path, weights_only=False, map_location="cpu")
        _calibrate(checkpoint["actor_state_dict"], "Actor")
        _calibrate(checkpoint["critic_state_dict"], "Critic")

        temporary_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as handle:
                temporary_path = Path(handle.name)
            torch.save(checkpoint, temporary_path)
            print(
                f"[{label} NORMALIZER] Calibrated final reference channels to "
                f"mean={list(reference_mean)}, "
                f"std={list(torch.sqrt(torch.tensor(reference_var)).tolist())}."
            )
            return original_load(
                self,
                str(temporary_path),
                load_cfg=load_cfg,
                strict=strict,
                map_location=map_location,
            )
        finally:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)

    OnPolicyRunner.load = load_with_pose_calibration


def _enable_resume_optimizer_lr_sync() -> None:
    """Keep PPO's adaptive-LR scalar synchronized with a resumed optimizer.

    RSL-RL restores the optimizer parameter-group learning rate but leaves
    ``PPO.learning_rate`` at the fresh config value. On the first adaptive-KL
    minibatch, that stale scalar is written back into the optimizer. For this
    task it turned a checkpoint LR near 5e-5 into 1e-3 and damaged the actor on
    every resume. Synchronizing the scalar makes a resume numerically faithful.
    """
    from rsl_rl.runners import OnPolicyRunner

    original_load = OnPolicyRunner.load

    def load_and_sync_lr(self, path, load_cfg=None, strict=True, map_location=None):
        infos = original_load(
            self,
            path,
            load_cfg=load_cfg,
            strict=strict,
            map_location=map_location,
        )
        optimizer_loaded = load_cfg is None or bool(load_cfg.get("optimizer"))
        if optimizer_loaded and self.alg.optimizer.param_groups:
            restored_lr = float(self.alg.optimizer.param_groups[0]["lr"])
            self.alg.learning_rate = restored_lr
            print(
                f"[RESUME LR] Synchronized adaptive PPO learning rate to "
                f"{restored_lr:.6g} from the checkpoint optimizer."
            )
        return infos

    OnPolicyRunner.load = load_and_sync_lr


def _enable_resume_burn_in(num_steps: int) -> None:
    """Flush fresh-scene reset transients before the first PPO update."""
    if not 1 <= num_steps <= 10000:
        raise ValueError("--resume-burn-in-steps must be between 1 and 10000")

    import torch
    from rsl_rl.runners import OnPolicyRunner

    original_learn = OnPolicyRunner.learn

    def learn_after_burn_in(self, num_learning_iterations, init_at_random_ep_len=False):
        print(
            f"[RESUME BURN-IN] Running {num_steps} deterministic policy steps "
            "before collecting PPO rollouts."
        )
        observations = self.env.get_observations().to(self.device)
        policy = self.get_inference_policy(device=self.device)
        with torch.inference_mode():
            for _ in range(num_steps):
                actions = policy(observations)
                observations, _, dones, _ = self.env.step(actions.to(self.env.device))
                observations = observations.to(self.device)
                policy.reset(dones)
        print(
            "[RESUME BURN-IN] Complete; starting PPO without randomized "
            "initial episode clocks."
        )
        return original_learn(
            self,
            num_learning_iterations=num_learning_iterations,
            init_at_random_ep_len=False,
        )

    OnPolicyRunner.learn = learn_after_burn_in


def _enable_actor_update_freeze(
    num_updates: int,
    post_freeze_learning_rate: float | None = None,
) -> None:
    """Train the critic first while preserving the loaded actor exactly."""
    if not 1 <= num_updates <= 10000:
        raise ValueError("--freeze-actor-updates must be between 1 and 10000")
    if (
        post_freeze_learning_rate is not None
        and not 1.0e-8 <= post_freeze_learning_rate <= 1.0
    ):
        raise ValueError(
            "--post-freeze-learning-rate must be between 1e-8 and 1"
        )

    from rsl_rl.algorithms import PPO

    original_update = PPO.update
    update_count = 0

    def update_with_actor_freeze(self):
        nonlocal update_count
        freeze_actor = update_count < num_updates
        if freeze_actor:
            self.actor.requires_grad_(False)
            if update_count == 0:
                print(
                    f"[ACTOR FREEZE] Preserving the loaded policy for "
                    f"{num_updates} PPO updates while the critic adapts."
                )
        try:
            result = original_update(self)
        finally:
            if freeze_actor:
                self.actor.requires_grad_(True)
        update_count += 1
        if update_count == num_updates:
            if post_freeze_learning_rate is not None:
                self.learning_rate = post_freeze_learning_rate
                for parameter_group in self.optimizer.param_groups:
                    parameter_group["lr"] = post_freeze_learning_rate
                print(
                    "[ACTOR FREEZE] Set post-warm-up actor/critic learning "
                    f"rate to {post_freeze_learning_rate:.6g}."
                )
            print(
                "[ACTOR FREEZE] Critic warm-up complete; actor updates are now enabled."
            )
        return result

    PPO.update = update_with_actor_freeze


def _enable_actor_update_scale(scale: float) -> None:
    """Scale the realized actor parameter step while leaving critic unchanged.

    Gradient scaling is ineffective for Adam's early updates because its
    normalization largely cancels a uniform gradient multiplier.  Interpolate
    the post-optimizer actor parameters toward their pre-update values instead;
    this bounds the actual policy change across all PPO epochs/minibatches.
    """
    if not 0.001 <= scale <= 1.0:
        raise ValueError("--actor-update-scale must be between 0.001 and 1")

    import torch
    from rsl_rl.algorithms import PPO

    original_update = PPO.update
    announced = False

    def update_with_scaled_actor_step(self):
        nonlocal announced
        actor_parameters = list(self.actor.parameters())
        before = [parameter.detach().clone() for parameter in actor_parameters]
        result = original_update(self)
        with torch.no_grad():
            for parameter, prior in zip(actor_parameters, before, strict=True):
                parameter.copy_(prior + scale * (parameter - prior))
        if not announced:
            print(
                f"[ACTOR UPDATE] Retaining {scale:g} of each realized policy "
                "parameter step; critic updates remain full strength."
            )
            announced = True
        return result

    PPO.update = update_with_scaled_actor_step


def _set_initial_command_level(level: float) -> None:
    """Pass a resume-time curriculum level into the environment config."""
    if not 0.1 <= level <= 1.0:
        raise ValueError("--initial-command-level must be between 0.1 and 1.0")
    os.environ["DROPBEAR_INITIAL_COMMAND_LEVEL"] = f"{level:g}"


def _set_env_float(name: str, value: float, *, minimum: float, maximum: float) -> None:
    """Validate and expose a wrapper-only environment or PPO override."""
    if not minimum <= value <= maximum:
        raise ValueError(f"{name} must be between {minimum:g} and {maximum:g}")
    os.environ[name] = f"{value:g}"


def _enable_live_training_preview(
    state_path: Path,
    source_env: int | None,
) -> None:
    """Publish one headless training environment for a separate renderer."""
    import gymnasium as gym

    original_make = gym.make

    def make_with_live_training_state(
        env_id: str,
        *args: Any,
        **kwargs: Any,
    ) -> gym.Env:
        env = original_make(env_id, *args, **kwargs)
        from dropbear_walk.live_training_bridge import LiveTrainingStateExporter

        return LiveTrainingStateExporter(
            env,
            state_path=state_path,
            source_env=source_env,
        )

    gym.make = make_with_live_training_state


warm_start_std = _pop_float_arg("--warm-start-std")
live_preview_state = _pop_string_arg("--live-preview-state")
live_preview_env = _pop_int_arg("--live-preview-env")
resume_burn_in_steps = _pop_int_arg("--resume-burn-in-steps")
freeze_actor_updates = _pop_int_arg("--freeze-actor-updates")
post_freeze_learning_rate = _pop_float_arg(
    "--post-freeze-learning-rate"
)
actor_update_scale = _pop_float_arg("--actor-update-scale")
rollout_steps = _pop_int_arg("--rollout-steps")
save_interval = _pop_int_arg("--save-interval")
initial_command_level = _pop_float_arg("--initial-command-level")
entropy_coef = _pop_float_arg("--entropy-coef")
learning_rate = _pop_float_arg("--learning-rate")
desired_kl = _pop_float_arg("--desired-kl")
ppo_schedule = _pop_string_arg("--ppo-schedule")
symmetry_mirror_loss_coeff = _pop_float_arg(
    "--symmetry-mirror-loss-coeff"
)
symmetry_data_augmentation = (
    "--symmetry-data-augmentation" in sys.argv
)
if symmetry_data_augmentation:
    sys.argv.remove("--symmetry-data-augmentation")
symmetry_mirror_loss = "--symmetry-mirror-loss" in sys.argv
if symmetry_mirror_loss:
    sys.argv.remove("--symmetry-mirror-loss")
termination_penalty = _pop_float_arg("--termination-penalty")
flat_orientation_weight = _pop_float_arg("--flat-orientation-weight")
torso_pendulum_weight = _pop_float_arg("--torso-pendulum-weight")
torso_pendulum_amplitude = _pop_float_arg("--torso-pendulum-amplitude")
torso_pendulum_std = _pop_float_arg("--torso-pendulum-std")
torso_roll_bias_horizon = _pop_float_arg("--torso-roll-bias-horizon")
torso_roll_bias_std = _pop_float_arg("--torso-roll-bias-std")
torso_pendulum_warmup = _pop_float_arg("--torso-pendulum-warmup")
tracking_std = _pop_float_arg("--tracking-std")
yaw_tracking_weight = _pop_float_arg("--yaw-tracking-weight")
yaw_tracking_std = _pop_float_arg("--yaw-tracking-std")
fixed_forward_speed = _pop_float_arg("--fixed-forward-speed")
forward_speed_range = _pop_float_pair_arg("--forward-speed-range")
fixed_lateral_speed = _pop_float_arg("--fixed-lateral-speed")
lateral_speed_range = _pop_float_pair_arg("--lateral-speed-range")
fixed_yaw_rate = _pop_float_arg("--fixed-yaw-rate")
yaw_rate_range = _pop_float_pair_arg("--yaw-rate-range")
planar_cardinal_commands = "--planar-cardinal-commands" in sys.argv
if planar_cardinal_commands:
    sys.argv.remove("--planar-cardinal-commands")
cardinal_forward_weight = _pop_float_arg("--cardinal-forward-weight")
cardinal_backward_weight = _pop_float_arg("--cardinal-backward-weight")
cardinal_left_weight = _pop_float_arg("--cardinal-left-weight")
cardinal_right_weight = _pop_float_arg("--cardinal-right-weight")
standing_env_fraction = _pop_float_arg("--standing-env-fraction")
stand_still_weight = _pop_float_arg("--stand-still-weight")
stand_velocity_weight = _pop_float_arg("--stand-velocity-weight")
stable_forward_weight = _pop_float_arg("--stable-forward-weight")
stable_forward_std = _pop_float_arg("--stable-forward-std")
com_stand_height = _pop_float_arg("--com-stand-height")
com_height_delta = _pop_float_arg("--com-height-delta")
com_height_error_scale = _pop_float_arg("--com-height-error-scale")
com_vertical_velocity_error_scale = _pop_float_arg(
    "--com-vertical-velocity-error-scale"
)
com_velocity_weight = _pop_float_arg("--com-velocity-weight")
com_planar_velocity_penalty_weight = _pop_float_arg(
    "--com-planar-velocity-penalty-weight"
)
com_planar_velocity_forward_scale = _pop_float_arg(
    "--com-planar-velocity-forward-scale"
)
com_planar_velocity_lateral_scale = _pop_float_arg(
    "--com-planar-velocity-lateral-scale"
)
com_height_weight = _pop_float_arg("--com-height-weight")
com_position_weight = _pop_float_arg("--com-position-weight")
com_velocity_xy_std = _pop_float_arg("--com-velocity-xy-std")
com_velocity_z_std = _pop_float_arg("--com-velocity-z-std")
com_height_std = _pop_float_arg("--com-height-std")
com_position_std = _pop_float_arg("--com-position-std")
gait_period = _pop_float_arg("--gait-period")
contact_timing_penalty_weight = _pop_float_arg(
    "--contact-timing-penalty-weight"
)
contact_timing_horizon = _pop_float_arg("--contact-timing-horizon")
contact_timing_warmup = _pop_float_arg("--contact-timing-warmup")
contact_min_air_time = _pop_float_arg("--contact-min-air-time")
contact_duty_std = _pop_float_arg("--contact-duty-std")
contact_rate_std = _pop_float_arg("--contact-rate-std")
contact_interval_std = _pop_float_arg("--contact-interval-std")
contact_flight_time_std = _pop_float_arg("--contact-flight-time-std")
arm_swing_weight = _pop_float_arg("--arm-swing-weight")
arm_swing_amplitude = _pop_float_arg("--arm-swing-amplitude")
arm_swing_std = _pop_float_arg("--arm-swing-std")
arm_extension_penalty_weight = _pop_float_arg(
    "--arm-extension-penalty-weight"
)
arm_extension_soft_limit = _pop_float_arg("--arm-extension-soft-limit")
arm_counterweight_penalty_weight = _pop_float_arg(
    "--arm-counterweight-penalty-weight"
)
arm_counterweight_soft_limit = _pop_float_arg(
    "--arm-counterweight-soft-limit"
)
foot_phase_velocity_weight = _pop_float_arg("--foot-phase-velocity-weight")
foot_phase_velocity_balance_mix = _pop_float_arg(
    "--foot-phase-velocity-balance-mix"
)
swing_foot_speed_factor = _pop_float_arg("--swing-foot-speed-factor")
swing_foot_forward_std = _pop_float_arg("--swing-foot-forward-std")
swing_foot_lateral_std = _pop_float_arg("--swing-foot-lateral-std")
alternating_knee_weight = _pop_float_arg("--alternating-knee-weight")
knee_reference_speed = _pop_float_arg("--knee-reference-speed")
knee_stance_offset = _pop_float_arg("--knee-stance-offset")
knee_swing_bend_offset = _pop_float_arg("--knee-swing-bend-offset")
knee_stance_std = _pop_float_arg("--knee-stance-std")
knee_swing_std = _pop_float_arg("--knee-swing-std")
knee_balance_mix = _pop_float_arg("--knee-balance-mix")
knee_swing_focus_mix = _pop_float_arg("--knee-swing-focus-mix")
knee_left_weight = _pop_float_arg("--knee-left-weight")
knee_right_weight = _pop_float_arg("--knee-right-weight")
alternating_step_through_weight = _pop_float_arg(
    "--alternating-step-through-weight"
)
bilateral_step_progress_weight = _pop_float_arg(
    "--bilateral-step-progress-weight"
)
bilateral_min_pass_distance = _pop_float_arg(
    "--bilateral-min-pass-distance"
)
bilateral_min_swing_bend = _pop_float_arg(
    "--bilateral-min-swing-bend"
)
step_reference_speed = _pop_float_arg("--step-reference-speed")
step_length_at_reference = _pop_float_arg("--step-length-at-reference")
step_length_std = _pop_float_arg("--step-length-std")
step_left_weight = _pop_float_arg("--step-left-weight")
step_right_weight = _pop_float_arg("--step-right-weight")
anticipatory_foot_placement_weight = _pop_float_arg(
    "--anticipatory-foot-placement-weight"
)
foot_nominal_forward_center = _pop_float_arg(
    "--foot-nominal-forward-center"
)
foot_nominal_lateral_center = _pop_float_arg(
    "--foot-nominal-lateral-center"
)
foot_nominal_half_width = _pop_float_arg("--foot-nominal-half-width")
foot_command_lead_time = _pop_float_arg("--foot-command-lead-time")
foot_placement_forward_std = _pop_float_arg(
    "--foot-placement-forward-std"
)
foot_placement_lateral_std = _pop_float_arg(
    "--foot-placement-lateral-std"
)
feet_self_interaction_penalty_weight = _pop_float_arg(
    "--feet-self-interaction-penalty-weight"
)
feet_min_distance = _pop_float_arg("--feet-min-distance")
feet_min_lateral_separation = _pop_float_arg(
    "--feet-min-lateral-separation"
)
feet_overlap_termination_distance = _pop_float_arg(
    "--feet-overlap-termination-distance"
)
feet_overlap_termination_lateral_separation = _pop_float_arg(
    "--feet-overlap-termination-lateral-separation"
)
feet_approach_penalty_weight = _pop_float_arg(
    "--feet-approach-penalty-weight"
)
feet_approach_distance = _pop_float_arg("--feet-approach-distance")
feet_approach_lateral_separation = _pop_float_arg(
    "--feet-approach-lateral-separation"
)
feet_closing_speed_scale = _pop_float_arg("--feet-closing-speed-scale")
feet_touchdown_clearance_weight = _pop_float_arg(
    "--feet-touchdown-clearance-weight"
)
feet_late_swing_phase_start = _pop_float_arg(
    "--feet-late-swing-phase-start"
)
feet_touchdown_distance = _pop_float_arg("--feet-touchdown-distance")
feet_touchdown_lateral_separation = _pop_float_arg(
    "--feet-touchdown-lateral-separation"
)
feet_touchdown_distance_std = _pop_float_arg(
    "--feet-touchdown-distance-std"
)
feet_touchdown_lateral_std = _pop_float_arg(
    "--feet-touchdown-lateral-std"
)
pose_tracking_weight = _pop_float_arg("--pose-tracking-weight")
pose_height_tracking_weight = _pop_float_arg("--pose-height-tracking-weight")
pose_trajectory_weight = _pop_float_arg("--pose-trajectory-weight")
pose_trajectory_velocity_weight = _pop_float_arg(
    "--pose-trajectory-velocity-weight"
)
pose_baseline_depth = _pop_float_arg("--pose-baseline-depth")
pose_depth_amplitude = _pop_float_arg("--pose-depth-amplitude")
pose_stand_height = _pop_float_arg("--pose-stand-height")
pose_crouch_height_delta = _pop_float_arg("--pose-crouch-height-delta")
pose_reset_ramp_s = _pop_float_arg("--pose-reset-ramp-s")
pose_residual_scale = _pop_float_arg("--pose-residual-scale")
pose_pg_pitch_offset = _pop_float_arg("--pose-pg-pitch-offset")
pose_knee_offset = _pop_float_arg("--pose-knee-offset")
pose_ankle67_offset = _pop_float_arg("--pose-ankle67-offset")
reset_joint_position_range = _pop_float_pair_arg("--reset-joint-position-range")
reset_joint_velocity_range = _pop_float_pair_arg("--reset-joint-velocity-range")
tracking_frame = _pop_string_arg("--tracking-frame")
actor_base_lin_vel = "--actor-base-lin-vel" in sys.argv
if actor_base_lin_vel:
    sys.argv.remove("--actor-base-lin-vel")
adapt_base_lin_vel_checkpoint = "--adapt-base-lin-vel-checkpoint" in sys.argv
if adapt_base_lin_vel_checkpoint:
    sys.argv.remove("--adapt-base-lin-vel-checkpoint")
adapt_pose_checkpoint = "--adapt-pose-checkpoint" in sys.argv
if adapt_pose_checkpoint:
    sys.argv.remove("--adapt-pose-checkpoint")
adapt_terrain_scan_checkpoint = (
    "--adapt-terrain-scan-checkpoint" in sys.argv
)
if adapt_terrain_scan_checkpoint:
    sys.argv.remove("--adapt-terrain-scan-checkpoint")
calibrate_pose_normalizer = "--calibrate-pose-normalizer" in sys.argv
if calibrate_pose_normalizer:
    sys.argv.remove("--calibrate-pose-normalizer")
calibrate_com_normalizer = "--calibrate-com-normalizer" in sys.argv
if calibrate_com_normalizer:
    sys.argv.remove("--calibrate-com-normalizer")
plane_only = "--plane-only" in sys.argv
if plane_only:
    sys.argv.remove("--plane-only")
obstacle_terrain = "--obstacle-terrain" in sys.argv
if obstacle_terrain:
    sys.argv.remove("--obstacle-terrain")
directional_obstacle_terrain = (
    "--directional-obstacle-terrain" in sys.argv
)
if directional_obstacle_terrain:
    sys.argv.remove("--directional-obstacle-terrain")
directional_obstacle_speed = _pop_float_arg(
    "--directional-obstacle-speed"
)
directional_turn_yaw_rate = _pop_float_arg(
    "--directional-turn-yaw-rate"
)
terrain_height_scan = "--terrain-height-scan" in sys.argv
if terrain_height_scan:
    sys.argv.remove("--terrain-height-scan")
push_forward_velocity = _pop_float_arg("--push-forward-velocity")
push_lateral_velocity = _pop_float_arg("--push-lateral-velocity")
push_interval_s = _pop_float_arg("--push-interval-s")
disable_pushes = "--disable-pushes" in sys.argv
if disable_pushes:
    sys.argv.remove("--disable-pushes")
reset_policy_joints_only = "--reset-policy-joints-only" in sys.argv
if reset_policy_joints_only:
    sys.argv.remove("--reset-policy-joints-only")
reciprocal_shoulder_actions = "--reciprocal-shoulder-actions" in sys.argv
if reciprocal_shoulder_actions:
    sys.argv.remove("--reciprocal-shoulder-actions")
shoulder_counterweight_scale = _pop_float_arg(
    "--shoulder-counterweight-scale"
)
gait_knee_action_adapter = "--gait-knee-action-adapter" in sys.argv
if gait_knee_action_adapter:
    sys.argv.remove("--gait-knee-action-adapter")
gait_knee_guide_strength = _pop_float_arg(
    "--gait-knee-guide-strength"
)
gait_mirror_action_adapter = "--gait-mirror-action-adapter" in sys.argv
if gait_mirror_action_adapter:
    sys.argv.remove("--gait-mirror-action-adapter")
gait_mirror_transfer_strength = _pop_float_arg(
    "--gait-mirror-transfer-strength"
)
pose_sequence = "--pose-sequence" in sys.argv
if pose_sequence:
    sys.argv.remove("--pose-sequence")
stagger_pose_phases = "--stagger-pose-phases" in sys.argv
if stagger_pose_phases:
    sys.argv.remove("--stagger-pose-phases")
pose_action_residual = "--pose-action-residual" in sys.argv
if pose_action_residual:
    sys.argv.remove("--pose-action-residual")
com_control = "--com-control" in sys.argv
if com_control:
    sys.argv.remove("--com-control")
    os.environ["DROPBEAR_COM_CONTROL"] = "1"
if initial_command_level is not None:
    _set_initial_command_level(initial_command_level)
if rollout_steps is not None:
    if not 8 <= rollout_steps <= 2000:
        raise ValueError("--rollout-steps must be between 8 and 2000")
    os.environ["DROPBEAR_ROLLOUT_STEPS"] = str(rollout_steps)
if save_interval is not None:
    if not 1 <= save_interval <= 10000:
        raise ValueError("--save-interval must be between 1 and 10000")
    os.environ["DROPBEAR_SAVE_INTERVAL"] = str(save_interval)
if entropy_coef is not None:
    _set_env_float("DROPBEAR_ENTROPY_COEF", entropy_coef, minimum=0.0, maximum=0.1)
if learning_rate is not None:
    _set_env_float(
        "DROPBEAR_LEARNING_RATE",
        learning_rate,
        minimum=1.0e-6,
        maximum=1.0e-2,
    )
if desired_kl is not None:
    _set_env_float(
        "DROPBEAR_DESIRED_KL",
        desired_kl,
        minimum=1.0e-5,
        maximum=1.0,
    )
if ppo_schedule is not None:
    if ppo_schedule not in {"adaptive", "fixed"}:
        raise ValueError("--ppo-schedule must be 'adaptive' or 'fixed'")
    os.environ["DROPBEAR_PPO_SCHEDULE"] = ppo_schedule
if symmetry_data_augmentation:
    os.environ["DROPBEAR_SYMMETRY_DATA_AUGMENTATION"] = "1"
if symmetry_mirror_loss:
    os.environ["DROPBEAR_SYMMETRY_MIRROR_LOSS"] = "1"
if symmetry_mirror_loss_coeff is not None:
    if not symmetry_mirror_loss:
        raise ValueError(
            "--symmetry-mirror-loss-coeff requires "
            "--symmetry-mirror-loss"
        )
    _set_env_float(
        "DROPBEAR_SYMMETRY_MIRROR_LOSS_COEFF",
        symmetry_mirror_loss_coeff,
        minimum=0.0,
        maximum=1000.0,
    )
if termination_penalty is not None:
    _set_env_float(
        "DROPBEAR_TERMINATION_PENALTY",
        termination_penalty,
        minimum=-1000.0,
        maximum=0.0,
    )
if flat_orientation_weight is not None:
    _set_env_float(
        "DROPBEAR_FLAT_ORIENTATION_WEIGHT",
        flat_orientation_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if torso_pendulum_weight is not None:
    _set_env_float(
        "DROPBEAR_TORSO_PENDULUM_WEIGHT",
        torso_pendulum_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if torso_pendulum_amplitude is not None:
    _set_env_float(
        "DROPBEAR_TORSO_PENDULUM_AMPLITUDE",
        torso_pendulum_amplitude,
        minimum=-0.5,
        maximum=0.5,
    )
if torso_pendulum_std is not None:
    _set_env_float(
        "DROPBEAR_TORSO_PENDULUM_STD",
        torso_pendulum_std,
        minimum=0.001,
        maximum=1.0,
    )
if torso_roll_bias_horizon is not None:
    _set_env_float(
        "DROPBEAR_TORSO_ROLL_BIAS_HORIZON",
        torso_roll_bias_horizon,
        minimum=0.1,
        maximum=60.0,
    )
if torso_roll_bias_std is not None:
    _set_env_float(
        "DROPBEAR_TORSO_ROLL_BIAS_STD",
        torso_roll_bias_std,
        minimum=0.001,
        maximum=1.0,
    )
if torso_pendulum_warmup is not None:
    _set_env_float(
        "DROPBEAR_TORSO_PENDULUM_WARMUP",
        torso_pendulum_warmup,
        minimum=0.0,
        maximum=60.0,
    )
if tracking_std is not None:
    _set_env_float(
        "DROPBEAR_TRACKING_STD",
        tracking_std,
        minimum=0.05,
        maximum=2.0,
    )
if yaw_tracking_weight is not None:
    _set_env_float(
        "DROPBEAR_YAW_TRACKING_WEIGHT",
        yaw_tracking_weight,
        minimum=0.0,
        maximum=100.0,
    )
if yaw_tracking_std is not None:
    _set_env_float(
        "DROPBEAR_YAW_TRACKING_STD",
        yaw_tracking_std,
        minimum=0.01,
        maximum=3.0,
    )
if fixed_forward_speed is not None:
    _set_env_float(
        "DROPBEAR_FIXED_FORWARD_SPEED",
        fixed_forward_speed,
        minimum=-2.0,
        maximum=2.0,
    )
if forward_speed_range is not None:
    if fixed_forward_speed is not None:
        raise ValueError("--forward-speed-range conflicts with --fixed-forward-speed")
    speed_min, speed_max = forward_speed_range
    if not -2.0 <= speed_min <= speed_max <= 2.0:
        raise ValueError("--forward-speed-range must satisfy -2 <= low <= high <= 2")
    os.environ["DROPBEAR_FORWARD_SPEED_MIN"] = f"{speed_min:g}"
    os.environ["DROPBEAR_FORWARD_SPEED_MAX"] = f"{speed_max:g}"
if fixed_lateral_speed is not None and lateral_speed_range is not None:
    raise ValueError(
        "--fixed-lateral-speed conflicts with --lateral-speed-range"
    )
if fixed_lateral_speed is not None:
    _set_env_float(
        "DROPBEAR_FIXED_LATERAL_SPEED",
        fixed_lateral_speed,
        minimum=-2.0,
        maximum=2.0,
    )
if lateral_speed_range is not None:
    lateral_min, lateral_max = lateral_speed_range
    if not -2.0 <= lateral_min <= lateral_max <= 2.0:
        raise ValueError(
            "--lateral-speed-range must satisfy -2 <= low <= high <= 2"
        )
    os.environ["DROPBEAR_LATERAL_SPEED_MIN"] = f"{lateral_min:g}"
    os.environ["DROPBEAR_LATERAL_SPEED_MAX"] = f"{lateral_max:g}"
if planar_cardinal_commands:
    if forward_speed_range is None or lateral_speed_range is None:
        raise ValueError(
            "--planar-cardinal-commands requires both --forward-speed-range "
            "and --lateral-speed-range"
        )
    forward_min, forward_max = forward_speed_range
    lateral_min, lateral_max = lateral_speed_range
    if not forward_min < 0.0 < forward_max:
        raise ValueError(
            "--planar-cardinal-commands requires a forward range spanning zero"
        )
    if not lateral_min < 0.0 < lateral_max:
        raise ValueError(
            "--planar-cardinal-commands requires a lateral range spanning zero"
        )
    os.environ["DROPBEAR_PLANAR_CARDINAL_COMMANDS"] = "1"
cardinal_direction_weights = (
    cardinal_forward_weight,
    cardinal_backward_weight,
    cardinal_left_weight,
    cardinal_right_weight,
)
if any(weight is not None for weight in cardinal_direction_weights):
    if not planar_cardinal_commands:
        raise ValueError(
            "Cardinal direction weights require --planar-cardinal-commands"
        )
    if any(weight is None for weight in cardinal_direction_weights):
        raise ValueError(
            "Specify all four cardinal direction weights together"
        )
    if sum(cardinal_direction_weights) <= 0.0:
        raise ValueError(
            "Cardinal direction weights must have a positive sum"
        )
    for environment_name, weight in zip(
        (
            "DROPBEAR_CARDINAL_FORWARD_WEIGHT",
            "DROPBEAR_CARDINAL_BACKWARD_WEIGHT",
            "DROPBEAR_CARDINAL_LEFT_WEIGHT",
            "DROPBEAR_CARDINAL_RIGHT_WEIGHT",
        ),
        cardinal_direction_weights,
        strict=True,
    ):
        _set_env_float(
            environment_name,
            weight,
            minimum=0.0,
            maximum=100.0,
        )
if disable_pushes and any(
    value is not None
    for value in (
        push_forward_velocity,
        push_lateral_velocity,
        push_interval_s,
    )
):
    raise ValueError(
        "--disable-pushes conflicts with explicit push configuration"
    )
if push_forward_velocity is not None:
    _set_env_float(
        "DROPBEAR_PUSH_FORWARD_VELOCITY",
        push_forward_velocity,
        minimum=0.0,
        maximum=2.0,
    )
if push_lateral_velocity is not None:
    _set_env_float(
        "DROPBEAR_PUSH_LATERAL_VELOCITY",
        push_lateral_velocity,
        minimum=0.0,
        maximum=2.0,
    )
if push_interval_s is not None:
    _set_env_float(
        "DROPBEAR_PUSH_INTERVAL_S",
        push_interval_s,
        minimum=0.5,
        maximum=60.0,
    )
if fixed_yaw_rate is not None and yaw_rate_range is not None:
    raise ValueError("--fixed-yaw-rate conflicts with --yaw-rate-range")
if fixed_yaw_rate is not None:
    _set_env_float(
        "DROPBEAR_FIXED_YAW_RATE",
        fixed_yaw_rate,
        minimum=-3.0,
        maximum=3.0,
    )
if yaw_rate_range is not None:
    yaw_min, yaw_max = yaw_rate_range
    if not -3.0 <= yaw_min <= yaw_max <= 3.0:
        raise ValueError("--yaw-rate-range must satisfy -3 <= low <= high <= 3")
    os.environ["DROPBEAR_YAW_RATE_MIN"] = f"{yaw_min:g}"
    os.environ["DROPBEAR_YAW_RATE_MAX"] = f"{yaw_max:g}"
if standing_env_fraction is not None:
    _set_env_float(
        "DROPBEAR_STANDING_ENV_FRACTION",
        standing_env_fraction,
        minimum=0.0,
        maximum=1.0,
    )
if stand_still_weight is not None:
    _set_env_float(
        "DROPBEAR_STAND_STILL_WEIGHT",
        stand_still_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if stand_velocity_weight is not None:
    _set_env_float(
        "DROPBEAR_STAND_VELOCITY_WEIGHT",
        stand_velocity_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if stable_forward_weight is not None:
    _set_env_float(
        "DROPBEAR_STABLE_FORWARD_WEIGHT",
        stable_forward_weight,
        minimum=0.0,
        maximum=100.0,
    )
if stable_forward_std is not None:
    _set_env_float(
        "DROPBEAR_STABLE_FORWARD_STD",
        stable_forward_std,
        minimum=0.01,
        maximum=2.0,
    )
if com_stand_height is not None:
    _set_env_float(
        "DROPBEAR_COM_STAND_HEIGHT",
        com_stand_height,
        minimum=0.0,
        maximum=3.0,
    )
if com_height_delta is not None:
    _set_env_float(
        "DROPBEAR_COM_HEIGHT_DELTA",
        com_height_delta,
        minimum=0.0,
        maximum=1.0,
    )
if com_height_error_scale is not None:
    _set_env_float(
        "DROPBEAR_COM_HEIGHT_ERROR_SCALE",
        com_height_error_scale,
        minimum=0.001,
        maximum=1.0,
    )
if com_vertical_velocity_error_scale is not None:
    _set_env_float(
        "DROPBEAR_COM_VERTICAL_VELOCITY_ERROR_SCALE",
        com_vertical_velocity_error_scale,
        minimum=0.001,
        maximum=5.0,
    )
if com_velocity_weight is not None:
    _set_env_float(
        "DROPBEAR_COM_VELOCITY_WEIGHT",
        com_velocity_weight,
        minimum=0.0,
        maximum=100.0,
    )
if com_planar_velocity_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_COM_PLANAR_VELOCITY_PENALTY_WEIGHT",
        com_planar_velocity_penalty_weight,
        minimum=-1000.0,
        maximum=0.0,
    )
if com_planar_velocity_forward_scale is not None:
    _set_env_float(
        "DROPBEAR_COM_PLANAR_VELOCITY_FORWARD_SCALE",
        com_planar_velocity_forward_scale,
        minimum=0.0,
        maximum=100.0,
    )
if com_planar_velocity_lateral_scale is not None:
    _set_env_float(
        "DROPBEAR_COM_PLANAR_VELOCITY_LATERAL_SCALE",
        com_planar_velocity_lateral_scale,
        minimum=0.0,
        maximum=100.0,
    )
if com_height_weight is not None:
    _set_env_float(
        "DROPBEAR_COM_HEIGHT_WEIGHT",
        com_height_weight,
        minimum=0.0,
        maximum=100.0,
    )
if com_position_weight is not None:
    _set_env_float(
        "DROPBEAR_COM_POSITION_WEIGHT",
        com_position_weight,
        minimum=0.0,
        maximum=100.0,
    )
if com_velocity_xy_std is not None:
    _set_env_float(
        "DROPBEAR_COM_VELOCITY_XY_STD",
        com_velocity_xy_std,
        minimum=0.001,
        maximum=5.0,
    )
if com_velocity_z_std is not None:
    _set_env_float(
        "DROPBEAR_COM_VELOCITY_Z_STD",
        com_velocity_z_std,
        minimum=0.001,
        maximum=5.0,
    )
if com_height_std is not None:
    _set_env_float(
        "DROPBEAR_COM_HEIGHT_STD",
        com_height_std,
        minimum=0.001,
        maximum=1.0,
    )
if com_position_std is not None:
    _set_env_float(
        "DROPBEAR_COM_POSITION_STD",
        com_position_std,
        minimum=0.001,
        maximum=10.0,
    )
if gait_period is not None:
    _set_env_float(
        "DROPBEAR_GAIT_PERIOD",
        gait_period,
        minimum=0.20,
        maximum=2.0,
    )
if contact_timing_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_TIMING_PENALTY_WEIGHT",
        contact_timing_penalty_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if contact_timing_horizon is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_TIMING_HORIZON",
        contact_timing_horizon,
        minimum=0.1,
        maximum=60.0,
    )
if contact_timing_warmup is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_TIMING_WARMUP",
        contact_timing_warmup,
        minimum=0.0,
        maximum=60.0,
    )
if contact_min_air_time is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_MIN_AIR_TIME",
        contact_min_air_time,
        minimum=0.0,
        maximum=2.0,
    )
if contact_duty_std is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_DUTY_STD",
        contact_duty_std,
        minimum=0.001,
        maximum=1.0,
    )
if contact_rate_std is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_RATE_STD",
        contact_rate_std,
        minimum=0.001,
        maximum=20.0,
    )
if contact_interval_std is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_INTERVAL_STD",
        contact_interval_std,
        minimum=0.001,
        maximum=10.0,
    )
if contact_flight_time_std is not None:
    _set_env_float(
        "DROPBEAR_CONTACT_FLIGHT_TIME_STD",
        contact_flight_time_std,
        minimum=0.001,
        maximum=10.0,
    )
if arm_swing_weight is not None:
    _set_env_float(
        "DROPBEAR_ARM_SWING_WEIGHT",
        arm_swing_weight,
        minimum=0.0,
        maximum=100.0,
    )
if arm_swing_amplitude is not None:
    _set_env_float(
        "DROPBEAR_ARM_SWING_AMPLITUDE",
        arm_swing_amplitude,
        minimum=0.0,
        maximum=1.5,
    )
if arm_swing_std is not None:
    _set_env_float(
        "DROPBEAR_ARM_SWING_STD",
        arm_swing_std,
        minimum=0.001,
        maximum=2.0,
    )
if arm_extension_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_ARM_EXTENSION_PENALTY_WEIGHT",
        arm_extension_penalty_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if arm_extension_soft_limit is not None:
    _set_env_float(
        "DROPBEAR_ARM_EXTENSION_SOFT_LIMIT",
        arm_extension_soft_limit,
        minimum=0.0,
        maximum=2.0,
    )
if arm_counterweight_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_ARM_COUNTERWEIGHT_PENALTY_WEIGHT",
        arm_counterweight_penalty_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if arm_counterweight_soft_limit is not None:
    _set_env_float(
        "DROPBEAR_ARM_COUNTERWEIGHT_SOFT_LIMIT",
        arm_counterweight_soft_limit,
        minimum=0.0,
        maximum=2.0,
    )
if foot_phase_velocity_weight is not None:
    _set_env_float(
        "DROPBEAR_FOOT_PHASE_VELOCITY_WEIGHT",
        foot_phase_velocity_weight,
        minimum=0.0,
        maximum=100.0,
    )
if foot_phase_velocity_balance_mix is not None:
    _set_env_float(
        "DROPBEAR_FOOT_PHASE_VELOCITY_BALANCE_MIX",
        foot_phase_velocity_balance_mix,
        minimum=0.0,
        maximum=1.0,
    )
if swing_foot_speed_factor is not None:
    _set_env_float(
        "DROPBEAR_SWING_FOOT_SPEED_FACTOR",
        swing_foot_speed_factor,
        minimum=0.1,
        maximum=10.0,
    )
if swing_foot_forward_std is not None:
    _set_env_float(
        "DROPBEAR_SWING_FOOT_FORWARD_STD",
        swing_foot_forward_std,
        minimum=0.001,
        maximum=5.0,
    )
if swing_foot_lateral_std is not None:
    _set_env_float(
        "DROPBEAR_SWING_FOOT_LATERAL_STD",
        swing_foot_lateral_std,
        minimum=0.001,
        maximum=5.0,
    )
if alternating_knee_weight is not None:
    _set_env_float(
        "DROPBEAR_ALTERNATING_KNEE_WEIGHT",
        alternating_knee_weight,
        minimum=0.0,
        maximum=100.0,
    )
if knee_reference_speed is not None:
    _set_env_float(
        "DROPBEAR_KNEE_REFERENCE_SPEED",
        knee_reference_speed,
        minimum=0.01,
        maximum=5.0,
    )
if knee_stance_offset is not None:
    _set_env_float(
        "DROPBEAR_KNEE_STANCE_OFFSET",
        knee_stance_offset,
        minimum=-2.0,
        maximum=2.0,
    )
if knee_swing_bend_offset is not None:
    _set_env_float(
        "DROPBEAR_KNEE_SWING_BEND_OFFSET",
        knee_swing_bend_offset,
        minimum=0.0,
        maximum=2.0,
    )
if knee_stance_std is not None:
    _set_env_float(
        "DROPBEAR_KNEE_STANCE_STD",
        knee_stance_std,
        minimum=0.001,
        maximum=2.0,
    )
if knee_swing_std is not None:
    _set_env_float(
        "DROPBEAR_KNEE_SWING_STD",
        knee_swing_std,
        minimum=0.001,
        maximum=2.0,
    )
if knee_balance_mix is not None:
    _set_env_float(
        "DROPBEAR_KNEE_BALANCE_MIX",
        knee_balance_mix,
        minimum=0.0,
        maximum=1.0,
    )
if knee_swing_focus_mix is not None:
    _set_env_float(
        "DROPBEAR_KNEE_SWING_FOCUS_MIX",
        knee_swing_focus_mix,
        minimum=0.0,
        maximum=1.0,
    )
if knee_left_weight is not None:
    _set_env_float(
        "DROPBEAR_KNEE_LEFT_WEIGHT",
        knee_left_weight,
        minimum=0.0,
        maximum=10.0,
    )
if knee_right_weight is not None:
    _set_env_float(
        "DROPBEAR_KNEE_RIGHT_WEIGHT",
        knee_right_weight,
        minimum=0.0,
        maximum=10.0,
    )
if alternating_step_through_weight is not None:
    _set_env_float(
        "DROPBEAR_ALTERNATING_STEP_THROUGH_WEIGHT",
        alternating_step_through_weight,
        minimum=0.0,
        maximum=100.0,
    )
if bilateral_step_progress_weight is not None:
    _set_env_float(
        "DROPBEAR_BILATERAL_STEP_PROGRESS_WEIGHT",
        bilateral_step_progress_weight,
        minimum=0.0,
        maximum=100.0,
    )
if bilateral_min_pass_distance is not None:
    _set_env_float(
        "DROPBEAR_BILATERAL_MIN_PASS_DISTANCE",
        bilateral_min_pass_distance,
        minimum=0.0,
        maximum=1.0,
    )
if bilateral_min_swing_bend is not None:
    _set_env_float(
        "DROPBEAR_BILATERAL_MIN_SWING_BEND",
        bilateral_min_swing_bend,
        minimum=0.0,
        maximum=2.0,
    )
if step_reference_speed is not None:
    _set_env_float(
        "DROPBEAR_STEP_REFERENCE_SPEED",
        step_reference_speed,
        minimum=0.01,
        maximum=5.0,
    )
if step_length_at_reference is not None:
    _set_env_float(
        "DROPBEAR_STEP_LENGTH_AT_REFERENCE",
        step_length_at_reference,
        minimum=0.01,
        maximum=2.0,
    )
if step_length_std is not None:
    _set_env_float(
        "DROPBEAR_STEP_LENGTH_STD",
        step_length_std,
        minimum=0.001,
        maximum=1.0,
    )
if step_left_weight is not None:
    _set_env_float(
        "DROPBEAR_STEP_LEFT_WEIGHT",
        step_left_weight,
        minimum=0.0,
        maximum=10.0,
    )
if step_right_weight is not None:
    _set_env_float(
        "DROPBEAR_STEP_RIGHT_WEIGHT",
        step_right_weight,
        minimum=0.0,
        maximum=10.0,
    )
if anticipatory_foot_placement_weight is not None:
    _set_env_float(
        "DROPBEAR_ANTICIPATORY_FOOT_PLACEMENT_WEIGHT",
        anticipatory_foot_placement_weight,
        minimum=0.0,
        maximum=100.0,
    )
if foot_nominal_forward_center is not None:
    _set_env_float(
        "DROPBEAR_FOOT_NOMINAL_FORWARD_CENTER",
        foot_nominal_forward_center,
        minimum=-2.0,
        maximum=2.0,
    )
if foot_nominal_lateral_center is not None:
    _set_env_float(
        "DROPBEAR_FOOT_NOMINAL_LATERAL_CENTER",
        foot_nominal_lateral_center,
        minimum=-2.0,
        maximum=2.0,
    )
if foot_nominal_half_width is not None:
    _set_env_float(
        "DROPBEAR_FOOT_NOMINAL_HALF_WIDTH",
        foot_nominal_half_width,
        minimum=0.001,
        maximum=1.0,
    )
if foot_command_lead_time is not None:
    _set_env_float(
        "DROPBEAR_FOOT_COMMAND_LEAD_TIME",
        foot_command_lead_time,
        minimum=0.0,
        maximum=2.0,
    )
if foot_placement_forward_std is not None:
    _set_env_float(
        "DROPBEAR_FOOT_PLACEMENT_FORWARD_STD",
        foot_placement_forward_std,
        minimum=0.001,
        maximum=2.0,
    )
if foot_placement_lateral_std is not None:
    _set_env_float(
        "DROPBEAR_FOOT_PLACEMENT_LATERAL_STD",
        foot_placement_lateral_std,
        minimum=0.001,
        maximum=2.0,
    )
if feet_self_interaction_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_FEET_SELF_INTERACTION_PENALTY_WEIGHT",
        feet_self_interaction_penalty_weight,
        minimum=-1000.0,
        maximum=0.0,
    )
if feet_min_distance is not None:
    _set_env_float(
        "DROPBEAR_FEET_MIN_DISTANCE",
        feet_min_distance,
        minimum=0.01,
        maximum=1.0,
    )
if feet_min_lateral_separation is not None:
    _set_env_float(
        "DROPBEAR_FEET_MIN_LATERAL_SEPARATION",
        feet_min_lateral_separation,
        minimum=0.01,
        maximum=1.0,
    )
if feet_overlap_termination_distance is not None:
    _set_env_float(
        "DROPBEAR_FEET_OVERLAP_TERMINATION_DISTANCE",
        feet_overlap_termination_distance,
        minimum=0.0,
        maximum=1.0,
    )
if feet_overlap_termination_lateral_separation is not None:
    _set_env_float(
        "DROPBEAR_FEET_OVERLAP_TERMINATION_LATERAL_SEPARATION",
        feet_overlap_termination_lateral_separation,
        minimum=0.0,
        maximum=1.0,
    )
if feet_approach_penalty_weight is not None:
    _set_env_float(
        "DROPBEAR_FEET_APPROACH_PENALTY_WEIGHT",
        feet_approach_penalty_weight,
        minimum=-1000.0,
        maximum=0.0,
    )
if feet_approach_distance is not None:
    _set_env_float(
        "DROPBEAR_FEET_APPROACH_DISTANCE",
        feet_approach_distance,
        minimum=0.01,
        maximum=1.0,
    )
if feet_approach_lateral_separation is not None:
    _set_env_float(
        "DROPBEAR_FEET_APPROACH_LATERAL_SEPARATION",
        feet_approach_lateral_separation,
        minimum=0.01,
        maximum=1.0,
    )
if feet_closing_speed_scale is not None:
    _set_env_float(
        "DROPBEAR_FEET_CLOSING_SPEED_SCALE",
        feet_closing_speed_scale,
        minimum=0.01,
        maximum=5.0,
    )
if feet_touchdown_clearance_weight is not None:
    _set_env_float(
        "DROPBEAR_FEET_TOUCHDOWN_CLEARANCE_WEIGHT",
        feet_touchdown_clearance_weight,
        minimum=0.0,
        maximum=1000.0,
    )
if feet_late_swing_phase_start is not None:
    _set_env_float(
        "DROPBEAR_FEET_LATE_SWING_PHASE_START",
        feet_late_swing_phase_start,
        minimum=0.56,
        maximum=0.99,
    )
if feet_touchdown_distance is not None:
    _set_env_float(
        "DROPBEAR_FEET_TOUCHDOWN_DISTANCE",
        feet_touchdown_distance,
        minimum=0.01,
        maximum=1.0,
    )
if feet_touchdown_lateral_separation is not None:
    _set_env_float(
        "DROPBEAR_FEET_TOUCHDOWN_LATERAL_SEPARATION",
        feet_touchdown_lateral_separation,
        minimum=0.01,
        maximum=1.0,
    )
if feet_touchdown_distance_std is not None:
    _set_env_float(
        "DROPBEAR_FEET_TOUCHDOWN_DISTANCE_STD",
        feet_touchdown_distance_std,
        minimum=0.001,
        maximum=1.0,
    )
if feet_touchdown_lateral_std is not None:
    _set_env_float(
        "DROPBEAR_FEET_TOUCHDOWN_LATERAL_STD",
        feet_touchdown_lateral_std,
        minimum=0.001,
        maximum=1.0,
    )
if pose_tracking_weight is not None:
    _set_env_float(
        "DROPBEAR_POSE_TRACKING_WEIGHT",
        pose_tracking_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if pose_height_tracking_weight is not None:
    _set_env_float(
        "DROPBEAR_POSE_HEIGHT_TRACKING_WEIGHT",
        pose_height_tracking_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if pose_trajectory_weight is not None:
    _set_env_float(
        "DROPBEAR_POSE_TRAJECTORY_WEIGHT",
        pose_trajectory_weight,
        minimum=-100.0,
        maximum=0.0,
    )
if pose_trajectory_velocity_weight is not None:
    _set_env_float(
        "DROPBEAR_POSE_TRAJECTORY_VELOCITY_WEIGHT",
        pose_trajectory_velocity_weight,
        minimum=0.0,
        maximum=100.0,
    )
if pose_baseline_depth is not None:
    _set_env_float(
        "DROPBEAR_POSE_BASELINE_DEPTH",
        pose_baseline_depth,
        minimum=-2.0,
        maximum=2.0,
    )
if pose_depth_amplitude is not None:
    _set_env_float(
        "DROPBEAR_POSE_DEPTH_AMPLITUDE",
        pose_depth_amplitude,
        minimum=0.0,
        maximum=2.0,
    )
if pose_stand_height is not None:
    _set_env_float(
        "DROPBEAR_POSE_STAND_HEIGHT",
        pose_stand_height,
        minimum=0.0,
        maximum=2.0,
    )
if pose_crouch_height_delta is not None:
    _set_env_float(
        "DROPBEAR_POSE_CROUCH_HEIGHT_DELTA",
        pose_crouch_height_delta,
        minimum=0.0,
        maximum=1.0,
    )
if pose_reset_ramp_s is not None:
    _set_env_float(
        "DROPBEAR_POSE_RESET_RAMP_S",
        pose_reset_ramp_s,
        minimum=0.0,
        maximum=16.0,
    )
if pose_residual_scale is not None:
    _set_env_float(
        "DROPBEAR_POSE_RESIDUAL_SCALE",
        pose_residual_scale,
        minimum=0.0,
        maximum=1.0,
    )
if pose_pg_pitch_offset is not None:
    _set_env_float(
        "DROPBEAR_POSE_PG_PITCH_OFFSET",
        pose_pg_pitch_offset,
        minimum=-2.0,
        maximum=2.0,
    )
if pose_knee_offset is not None:
    _set_env_float(
        "DROPBEAR_POSE_KNEE_OFFSET",
        pose_knee_offset,
        minimum=-2.0,
        maximum=2.0,
    )
if pose_ankle67_offset is not None:
    _set_env_float(
        "DROPBEAR_POSE_ANKLE67_OFFSET",
        pose_ankle67_offset,
        minimum=-2.0,
        maximum=2.0,
    )
if reset_joint_position_range is not None:
    reset_min, reset_max = reset_joint_position_range
    if not 0.1 <= reset_min <= reset_max <= 2.0:
        raise ValueError(
            "--reset-joint-position-range must satisfy "
            "0.1 <= low <= high <= 2.0"
        )
    os.environ["DROPBEAR_RESET_JOINT_POSITION_MIN"] = f"{reset_min:g}"
    os.environ["DROPBEAR_RESET_JOINT_POSITION_MAX"] = f"{reset_max:g}"
if reset_joint_velocity_range is not None:
    reset_vel_min, reset_vel_max = reset_joint_velocity_range
    if not -5.0 <= reset_vel_min <= reset_vel_max <= 5.0:
        raise ValueError(
            "--reset-joint-velocity-range must satisfy "
            "-5.0 <= low <= high <= 5.0"
        )
    os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MIN"] = f"{reset_vel_min:g}"
    os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MAX"] = f"{reset_vel_max:g}"
if tracking_frame is not None:
    if tracking_frame not in {"yaw", "body"}:
        raise ValueError("--tracking-frame must be either 'yaw' or 'body'")
    os.environ["DROPBEAR_TRACKING_FRAME"] = tracking_frame
if actor_base_lin_vel:
    os.environ["DROPBEAR_ACTOR_BASE_LIN_VEL"] = "1"
if adapt_base_lin_vel_checkpoint:
    if not actor_base_lin_vel:
        raise ValueError(
            "--adapt-base-lin-vel-checkpoint requires --actor-base-lin-vel"
        )
    if warm_start_std is None:
        raise ValueError(
            "--adapt-base-lin-vel-checkpoint requires --warm-start-std "
            "because the old optimizer moments are shape-incompatible"
        )
if obstacle_terrain and directional_obstacle_terrain:
    raise ValueError(
        "--obstacle-terrain conflicts with "
        "--directional-obstacle-terrain"
    )
if directional_obstacle_terrain:
    if plane_only:
        raise ValueError(
            "--directional-obstacle-terrain conflicts with --plane-only"
        )
    speed = (
        0.20
        if directional_obstacle_speed is None
        else directional_obstacle_speed
    )
    turn_rate = (
        0.20
        if directional_turn_yaw_rate is None
        else directional_turn_yaw_rate
    )
    if not 0.05 <= speed <= 1.0:
        raise ValueError(
            "--directional-obstacle-speed must be within [0.05, 1.0]"
        )
    if not 0.05 <= turn_rate <= 1.5:
        raise ValueError(
            "--directional-turn-yaw-rate must be within [0.05, 1.5]"
        )
    os.environ["DROPBEAR_DIRECTIONAL_OBSTACLE_TERRAIN"] = "1"
    os.environ["DROPBEAR_DIRECTIONAL_OBSTACLE_COMMANDS"] = "1"
    os.environ["DROPBEAR_DIRECTIONAL_OBSTACLE_SPEED"] = f"{speed:g}"
    os.environ["DROPBEAR_DIRECTIONAL_TURN_YAW_RATE"] = f"{turn_rate:g}"
    os.environ["DROPBEAR_TERRAIN_HEIGHT_SCAN"] = "1"
elif obstacle_terrain:
    if plane_only:
        raise ValueError("--obstacle-terrain conflicts with --plane-only")
    os.environ["DROPBEAR_OBSTACLE_TERRAIN"] = "1"
    os.environ["DROPBEAR_TERRAIN_HEIGHT_SCAN"] = "1"
elif terrain_height_scan:
    os.environ["DROPBEAR_TERRAIN_HEIGHT_SCAN"] = "1"
if not directional_obstacle_terrain:
    if directional_obstacle_speed is not None:
        raise ValueError(
            "--directional-obstacle-speed requires "
            "--directional-obstacle-terrain"
        )
    if directional_turn_yaw_rate is not None:
        raise ValueError(
            "--directional-turn-yaw-rate requires "
            "--directional-obstacle-terrain"
        )
if adapt_terrain_scan_checkpoint:
    if not (
        obstacle_terrain
        or directional_obstacle_terrain
        or terrain_height_scan
    ):
        raise ValueError(
            "--adapt-terrain-scan-checkpoint requires "
            "--obstacle-terrain, --directional-obstacle-terrain, "
            "or --terrain-height-scan"
        )
    if warm_start_std is None:
        raise ValueError(
            "--adapt-terrain-scan-checkpoint requires --warm-start-std "
            "because the old optimizer moments are shape-incompatible"
        )
if plane_only:
    os.environ["DROPBEAR_PLANE_ONLY"] = "1"
if disable_pushes:
    os.environ["DROPBEAR_DISABLE_PUSHES"] = "1"
if reset_policy_joints_only:
    os.environ["DROPBEAR_RESET_POLICY_JOINTS_ONLY"] = "1"
if reciprocal_shoulder_actions:
    os.environ["DROPBEAR_RECIPROCAL_SHOULDER_ACTIONS"] = "1"
    if shoulder_counterweight_scale is not None:
        raise ValueError(
            "--reciprocal-shoulder-actions conflicts with "
            "--shoulder-counterweight-scale"
        )
if shoulder_counterweight_scale is not None:
    _set_env_float(
        "DROPBEAR_SHOULDER_COUNTERWEIGHT_SCALE",
        shoulder_counterweight_scale,
        minimum=0.0,
        maximum=1.0,
    )
    os.environ["DROPBEAR_RECIPROCAL_SHOULDER_ACTIONS"] = "1"
if gait_knee_action_adapter:
    if reciprocal_shoulder_actions or shoulder_counterweight_scale is not None:
        raise ValueError(
            "--gait-knee-action-adapter currently conflicts with the "
            "reciprocal shoulder action adapter"
        )
    os.environ["DROPBEAR_GAIT_KNEE_ACTION_ADAPTER"] = "1"
    if gait_knee_guide_strength is None:
        os.environ["DROPBEAR_GAIT_KNEE_GUIDE_STRENGTH"] = "0.25"
elif gait_knee_guide_strength is not None:
    raise ValueError(
        "--gait-knee-guide-strength requires --gait-knee-action-adapter"
    )
if gait_knee_guide_strength is not None:
    _set_env_float(
        "DROPBEAR_GAIT_KNEE_GUIDE_STRENGTH",
        gait_knee_guide_strength,
        minimum=0.0,
        maximum=1.0,
    )
if gait_mirror_action_adapter:
    if (
        gait_knee_action_adapter
        or reciprocal_shoulder_actions
        or shoulder_counterweight_scale is not None
    ):
        raise ValueError(
            "--gait-mirror-action-adapter conflicts with the other "
            "walking action adapters"
        )
    os.environ["DROPBEAR_GAIT_MIRROR_ACTION_ADAPTER"] = "1"
    if gait_mirror_transfer_strength is None:
        os.environ["DROPBEAR_GAIT_MIRROR_TRANSFER_STRENGTH"] = "0.20"
elif gait_mirror_transfer_strength is not None:
    raise ValueError(
        "--gait-mirror-transfer-strength requires "
        "--gait-mirror-action-adapter"
    )
if gait_mirror_transfer_strength is not None:
    _set_env_float(
        "DROPBEAR_GAIT_MIRROR_TRANSFER_STRENGTH",
        gait_mirror_transfer_strength,
        minimum=0.0,
        maximum=1.0,
    )
if pose_sequence:
    os.environ["DROPBEAR_POSE_SEQUENCE"] = "1"
elif pose_reset_ramp_s is not None:
    raise ValueError("--pose-reset-ramp-s requires --pose-sequence")
if stagger_pose_phases:
    if not pose_sequence:
        raise ValueError("--stagger-pose-phases requires --pose-sequence")
    os.environ["DROPBEAR_POSE_PHASE_STAGGER"] = "1"
if pose_action_residual:
    if not pose_sequence:
        raise ValueError("--pose-action-residual requires --pose-sequence")
    os.environ["DROPBEAR_POSE_ACTION_RESIDUAL"] = "1"
    if pose_residual_scale is None:
        os.environ["DROPBEAR_POSE_RESIDUAL_SCALE"] = "1"
elif pose_residual_scale is not None:
    raise ValueError("--pose-residual-scale requires --pose-action-residual")
if adapt_pose_checkpoint:
    if not pose_sequence:
        raise ValueError("--adapt-pose-checkpoint requires --pose-sequence")
    if warm_start_std is None:
        raise ValueError(
            "--adapt-pose-checkpoint requires --warm-start-std because the "
            "old optimizer moments are shape-incompatible"
        )
if calibrate_pose_normalizer and not pose_sequence:
    raise ValueError("--calibrate-pose-normalizer requires --pose-sequence")
if calibrate_com_normalizer and not com_control:
    raise ValueError("--calibrate-com-normalizer requires --com-control")
if calibrate_pose_normalizer and calibrate_com_normalizer:
    raise ValueError(
        "--calibrate-pose-normalizer conflicts with --calibrate-com-normalizer"
    )

# Register tasks only after wrapper-only environment settings are available.
import dropbear_walk  # noqa: E402,F401

# RSL-RL's checkpoint loader needs this correction for every faithful optimizer
# resume. Actor-only warm starts wrap it below and deliberately skip the sync
# because their optimizer is freshly initialized.
_enable_resume_optimizer_lr_sync()
if warm_start_std is not None:
    _enable_exploration_warm_start(warm_start_std)
if calibrate_pose_normalizer:
    _enable_pose_reference_normalizer_calibration()
if calibrate_com_normalizer:
    _enable_pose_reference_normalizer_calibration(
        reference_mean=(0.0, 0.0),
        reference_var=(1.0, 1.0),
        label="COM",
    )
if adapt_base_lin_vel_checkpoint:
    _enable_actor_base_lin_vel_checkpoint_adapter()
if adapt_pose_checkpoint:
    _enable_pose_reference_checkpoint_adapter()
if adapt_terrain_scan_checkpoint:
    _enable_appended_observation_checkpoint_adapter("terrain-height scan")
if resume_burn_in_steps is not None:
    _enable_resume_burn_in(resume_burn_in_steps)
if freeze_actor_updates is not None:
    _enable_actor_update_freeze(
        freeze_actor_updates,
        post_freeze_learning_rate=post_freeze_learning_rate,
    )
elif post_freeze_learning_rate is not None:
    raise ValueError(
        "--post-freeze-learning-rate requires --freeze-actor-updates"
    )
if actor_update_scale is not None:
    _enable_actor_update_scale(actor_update_scale)
if live_preview_env is not None and live_preview_state is None:
    raise ValueError("--live-preview-env requires --live-preview-state")
if (
    live_preview_state is not None
    and int(os.environ.get("RANK", "0")) == 0
):
    _enable_live_training_preview(
        Path(live_preview_state),
        live_preview_env,
    )

if not TRAIN_SCRIPT.is_file():
    raise FileNotFoundError(f"Isaac Lab trainer not found: {TRAIN_SCRIPT}")

# Match Python's normal script execution semantics so the trainer's local
# ``cli_args`` module resolves without modifying the Isaac Lab checkout.
sys.path.insert(0, str(TRAIN_SCRIPT.parent))
runpy.run_path(str(TRAIN_SCRIPT), run_name="__main__")