"""S10 terrain v4 + upward recovery task. This is a separate task, not an inheritance of v4 or v7. It keeps the v4 53-value Actor frame, terrain revision, command schedule, terrain proportions, rolling rewards and replay, and changes only: * the existing ``upright`` reward into a sign-aware upward reward (no new term); * clamps on the large raw penalty terms; * random roll/pitch/yaw at reset; * a 4-second sustained bad-orientation termination. Reference: IsaacLab S10 ``upward_env_cfg.py`` for the upward/recovery state machine only. Its speed ranges, reward set and terrain config are not copied. """ import math from .base import ActorSpec, TaskSpec, register ACTOR = ActorSpec( name="s10-v4-upward-53", terms=(3, 3, 12, 12, 4, 16, 3), frame_terms=( {"name": "base_ang_vel", "dim": 3, "scale": 1.0}, {"name": "projected_gravity", "dim": 3, "scale": 1.0}, {"name": "joint_pos_rel", "dim": 12, "scale": 1.0}, {"name": "joint_vel_rel", "dim": 12, "scale": 0.05}, {"name": "wheel_vel_rel", "dim": 4, "scale": 0.05}, {"name": "last_action", "dim": 16, "scale": 1.0}, {"name": "command", "dim": 3, "scale": 1.0}, ), wheel_vel=True, ) def env_cfg(play: bool, history: int, seed: int): from mjlab.managers import RewardTermCfg, TerminationTermCfg from mjlab.managers.scene_entity_config import SceneEntityCfg from s10_supermimic.core.upward import ( BAD_ORIENTATION_WEIGHT, UPWARD_REWARD_CLAMPS, BadOrientationSustained, bad_orientation_penalty, clamped_action_rate, clamped_dof_pos_limits, clamped_joint_torques, clamped_leg_posture, clamped_rolling_cost, clamped_standstill, reset_root_state_upward, upright_upward, upward_curriculum, ) from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg cfg = s10_velocity_terrain_env_cfg( play=play, history=history, seed=seed, revision=4, actor_wheel_vel=True) # A 10 s episode lets each env complete 20 episodes per curriculum step # within a 20k-iteration run. The 4 s bad-orientation termination still # leaves 6 s after a successful recovery. Play/evaluation keeps the normal # terrain-course horizon. if not play: cfg.episode_length_s = 10.0 # Random orientation at reset. Training starts at +/-60 degrees and advances # per environment through upward_curriculum(); play/eval uses the full sphere # so the checkpoint is tested on the final distribution. tilt = math.pi if play else math.radians(60.0) pose = cfg.events["reset_base"].params["pose_range"] pose.update( z=(0.0, 0.2), roll=(-tilt, tilt), pitch=(-tilt, tilt), yaw=(-math.pi, math.pi), ) cfg.events["reset_base"].func = reset_root_state_upward # Existing upright term becomes the normalized upward reward. Max rate 2.0, # i.e. 0.04 per 20 ms step, the same ratio to tracking (2.0 each) as IsaacLab. upright = cfg.rewards["upright"] upright_asset = upright.params.get( "asset_cfg", SceneEntityCfg("robot", body_names=("base_link",))) upright.func = upright_upward upright.weight = 2.0 upright.params = {"asset_cfg": upright_asset} # IsaacLab-style recovery shaping: no reward for termination itself, but a # continuous per-step penalty while the orientation is bad. Weight -40 keeps # the same 20x bad/upward ratio as IsaacLab (-80 vs max 4; here -40 vs max 2). cfg.rewards["terminated"].weight = 0.0 cfg.rewards["bad_orientation_penalty"] = RewardTermCfg( func=bad_orientation_penalty, weight=BAD_ORIENTATION_WEIGHT, params={"asset_cfg": SceneEntityCfg("robot")}, ) # Clamp the large raw penalties. The weights and the rest of the v4 reward # set are unchanged; the caps only prevent recovery poses from dominating. leg = cfg.rewards["leg_posture"] leg.params = {"asset_cfg": leg.params["asset_cfg"], "max_value": UPWARD_REWARD_CLAMPS["leg_posture"]} leg.func = clamped_leg_posture standstill = cfg.rewards["standstill"] standstill.params = {"wheel_cfg": standstill.params["wheel_cfg"], "max_value": UPWARD_REWARD_CLAMPS["standstill"]} standstill.func = clamped_standstill cfg.rewards["action_rate_l2"].func = clamped_action_rate cfg.rewards["action_rate_l2"].params = { "max_value": UPWARD_REWARD_CLAMPS["action_rate_l2"]} cfg.rewards["joint_torques_l2"].func = clamped_joint_torques cfg.rewards["joint_torques_l2"].params = { "max_value": UPWARD_REWARD_CLAMPS["joint_torques_l2"]} limits = cfg.rewards["dof_pos_limits"] limits.params = {"asset_cfg": limits.params["asset_cfg"], "max_value": UPWARD_REWARD_CLAMPS["dof_pos_limits"]} limits.func = clamped_dof_pos_limits for name in ("unnecessary_wheel_lift", "rolling_leg_motion"): term = cfg.rewards[name] term.func = clamped_rolling_cost term.params = {"term": name, "max_value": UPWARD_REWARD_CLAMPS[name]} # Replace v4's immediate 70-degree fall with the upward recovery state # machine: bad orientation is tolerated for 4 s, then the episode ends. cfg.terminations["fell_over"] = TerminationTermCfg( func=BadOrientationSustained, params={"duration_s": 4.0, "asset_cfg": SceneEntityCfg("robot")}, ) # Each environment advances its own terrain/tilt level from its last 20 # recovery outcomes (16/20 promotes); there is no global batch gate. if not play: cfg.curriculum["course"].func = upward_curriculum return cfg TASK = register(TaskSpec( name="s10-v4-upward", description="S10 v4 recovery task: upward reward, monotonic tilt severity, per-env 20-episode curriculum, horizon 64", terrain_revision=4, actor=ACTOR, env_cfg=env_cfg, defaults={ "terrain": True, "terrain_revision": 4, "num_envs": 4096, "iterations": 30000, "num_steps_per_env": 1, "buffer_len": 512, "wm_window": 8, "wm_batch": 2048, "wm_updates": 4, "imag_batch": 2048, "horizon": 64, "warmup": 50, "action_noise": 0.22, "critic_epochs": 5, "critic_minibatches": 4, "save_interval": 500, "log_interval": 10, "command_curriculum": True, "command_curriculum_start": 0, "prioritized_replay": True, "archive_size": 100000, "actor_ramp_iterations": 1000, "actor_hold_iterations": 0, "lr_decay_iterations": 30000, "lr_final_frac": 0.03, "history": 5, "wm_hidden": [512, 512, 512], "actor_hidden": [256, 128, 64], "critic_hidden": [256, 128, 64], "compile_imagination": False, }, ))