173 lines
6.3 KiB
Python
173 lines
6.3 KiB
Python
"""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,
|
|
},
|
|
))
|