#!/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__")