[software] 添加16DOF早期训练仿真与Sim2Real闭环
This commit is contained in:
@@ -0,0 +1,25 @@
|
||||
from mjlab.tasks.registry import register_mjlab_task
|
||||
from .config.env_cfgs import flat_env_cfg, rough_env_cfg, crawl_env_cfg
|
||||
from .config.rl_cfg import flat_ppo_runner_cfg, rough_ppo_runner_cfg, crawl_ppo_runner_cfg
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Robot-Flat-v0",
|
||||
env_cfg=flat_env_cfg(),
|
||||
play_env_cfg=flat_env_cfg(play=True),
|
||||
rl_cfg=flat_ppo_runner_cfg(),
|
||||
)
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Robot-Rough-v0",
|
||||
env_cfg=rough_env_cfg(),
|
||||
play_env_cfg=rough_env_cfg(play=True),
|
||||
rl_cfg=rough_ppo_runner_cfg(),
|
||||
)
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Robot-Crawl-v0",
|
||||
env_cfg=crawl_env_cfg(),
|
||||
play_env_cfg=crawl_env_cfg(play=True),
|
||||
rl_cfg=crawl_ppo_runner_cfg(),
|
||||
)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,685 @@
|
||||
"""Robot Environment Configurations for RL Locomotion Training.
|
||||
|
||||
This module defines the Manager-Based RL Environment Configurations for unitree robots
|
||||
equipped with actuated wheels and leg joints. It handles sensors, actuators, command
|
||||
generators, observation/critic terms, event randomizations, rewards, and terminations
|
||||
for flat ground, rough terrains, and crawling tasks.
|
||||
"""
|
||||
|
||||
import math
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
from mjlab.envs import mdp as envs_mdp
|
||||
from mjlab.envs.mdp import dr as envs_dr
|
||||
from mjlab.sim import SimulationCfg, MujocoCfg
|
||||
from mjlab.managers.action_manager import ActionTermCfg
|
||||
from mjlab.managers.command_manager import CommandTermCfg
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.metrics_manager import MetricsTermCfg
|
||||
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.sensor import (
|
||||
ContactMatch,
|
||||
ContactSensorCfg,
|
||||
ObjRef,
|
||||
RayCastSensorCfg,
|
||||
GridPatternCfg,
|
||||
)
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
from mjlab.tasks.velocity import mdp as velocity_mdp
|
||||
from mjlab.terrains import (
|
||||
TerrainEntityCfg,
|
||||
TerrainGeneratorCfg,
|
||||
BoxFlatTerrainCfg,
|
||||
BoxPyramidStairsTerrainCfg,
|
||||
BoxInvertedPyramidStairsTerrainCfg,
|
||||
BoxRandomGridTerrainCfg,
|
||||
HfRandomUniformTerrainCfg,
|
||||
HfPerlinNoiseTerrainCfg,
|
||||
HfPyramidSlopedTerrainCfg,
|
||||
)
|
||||
from ..terrains import RCWallTerrainCfg, RCLowBarTerrainCfg
|
||||
from mjlab.utils.noise import UniformNoiseCfg as Unoise
|
||||
from mjlab.viewer import ViewerConfig
|
||||
|
||||
from ..robot_cfg import get_robot_cfg, get_robot_crawl_cfg, LEG_POS_SCALE, WHEEL_VEL_SCALE
|
||||
from ..mdp.lowpass_actions import JointPositionDelayedLowPassActionCfg, JointVelocityDelayedLowPassActionCfg
|
||||
from ..mdp.disturbances import apply_continuous_disturbance
|
||||
from ..mdp.only_positive_rewards import enable_only_positive_rewards
|
||||
from ..mdp.rewards import (
|
||||
track_linear_velocity,
|
||||
track_linear_velocity_l1,
|
||||
track_angular_velocity,
|
||||
base_height_l2,
|
||||
safe_base_lin_vel,
|
||||
safe_foot_contact,
|
||||
safe_height_scan,
|
||||
wheel_roll_tracking,
|
||||
adaptive_leg_motion_penalty,
|
||||
leg_symmetry,
|
||||
contact_fraction_reward,
|
||||
stand_still,
|
||||
hip_deviation,
|
||||
joint_deviation_l2,
|
||||
flat_orientation_l2,
|
||||
lin_vel_z_l2,
|
||||
crawl_height_reward,
|
||||
terrain_level_bonus,
|
||||
action_rate_curriculum_l2,
|
||||
variable_posture,
|
||||
joint_pos_penalty,
|
||||
joint_mirror,
|
||||
feet_contact_without_cmd,
|
||||
upright_roll_only,
|
||||
pitch_control_penalty,
|
||||
)
|
||||
from ..mdp.curriculums import terrain_levels_vel_strict
|
||||
from ..mdp.commands import UniformThresholdVelocityCommandCfg
|
||||
|
||||
# Constant Definitions
|
||||
WHEEL_NAMES = ("fl", "fr", "rl", "rr")
|
||||
|
||||
|
||||
def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
"""Create the base environment configuration containing sensors, commands, and default policies."""
|
||||
|
||||
# ------------------
|
||||
# Sensors Definition
|
||||
# ------------------
|
||||
feet_ground_cfg = ContactSensorCfg(
|
||||
name="feet_ground_contact",
|
||||
primary=ContactMatch(
|
||||
mode="body",
|
||||
pattern=tuple(f"{w}_wheel_Link" for w in WHEEL_NAMES),
|
||||
entity="robot",
|
||||
),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
reduce="none",
|
||||
num_slots=1,
|
||||
history_length=3,
|
||||
track_air_time=True,
|
||||
)
|
||||
|
||||
base_ground_cfg = ContactSensorCfg(
|
||||
name="base_ground_contact",
|
||||
primary=ContactMatch(mode="body", pattern="base_link", entity="robot"),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found",),
|
||||
reduce="none",
|
||||
num_slots=1,
|
||||
history_length=4,
|
||||
)
|
||||
|
||||
body_collision_cfg = ContactSensorCfg(
|
||||
name="body_collision",
|
||||
primary=ContactMatch(
|
||||
mode="body",
|
||||
pattern=(".*_hip_abduction_Link", ".*_hip_pitch_Link", ".*_knee_Link"),
|
||||
entity="robot",
|
||||
),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
reduce="none",
|
||||
num_slots=1,
|
||||
history_length=4,
|
||||
)
|
||||
|
||||
terrain_scan = RayCastSensorCfg(
|
||||
name="height_scanner",
|
||||
frame=ObjRef(type="body", name="base_link", entity="robot"),
|
||||
pattern=GridPatternCfg(resolution=0.08, size=(1.6, 1.0)),
|
||||
ray_alignment="yaw",
|
||||
max_distance=5.0,
|
||||
exclude_parent_body=True,
|
||||
include_geom_groups=(0,),
|
||||
debug_vis=False,
|
||||
)
|
||||
|
||||
# ------------------
|
||||
# Observations Setup
|
||||
# ------------------
|
||||
actor_terms = {
|
||||
"base_ang_vel": ObservationTermCfg(
|
||||
func=envs_mdp.base_ang_vel, scale=0.25,
|
||||
noise=Unoise(n_min=-0.2, n_max=0.2),
|
||||
),
|
||||
"projected_gravity": ObservationTermCfg(
|
||||
func=velocity_mdp.projected_gravity,
|
||||
noise=Unoise(n_min=-0.05, n_max=0.05),
|
||||
),
|
||||
"command": ObservationTermCfg(
|
||||
func=velocity_mdp.generated_commands,
|
||||
params={"command_name": "twist"},
|
||||
),
|
||||
"joint_pos": ObservationTermCfg(
|
||||
func=envs_mdp.joint_pos_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(
|
||||
".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint",
|
||||
))},
|
||||
noise=Unoise(n_min=-0.01, n_max=0.01),
|
||||
),
|
||||
"joint_vel": ObservationTermCfg(
|
||||
func=envs_mdp.joint_vel_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(
|
||||
".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint",
|
||||
))},
|
||||
scale=0.05, noise=Unoise(n_min=-1.5, n_max=1.5),
|
||||
),
|
||||
"wheel_vel": ObservationTermCfg(
|
||||
func=envs_mdp.joint_vel_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))},
|
||||
scale=0.05, noise=Unoise(n_min=-1.0, n_max=1.0),
|
||||
),
|
||||
"actions": ObservationTermCfg(func=velocity_mdp.last_action, history_length=1),
|
||||
}
|
||||
|
||||
critic_terms = {
|
||||
**actor_terms,
|
||||
"base_lin_vel": ObservationTermCfg(func=safe_base_lin_vel, scale=2.0),
|
||||
"foot_contact": ObservationTermCfg(
|
||||
func=safe_foot_contact, params={"sensor_name": "feet_ground_contact"},
|
||||
),
|
||||
"height_scan": ObservationTermCfg(
|
||||
func=safe_height_scan, params={"sensor_name": "height_scanner"},
|
||||
clip=(-1.0, 1.0),
|
||||
),
|
||||
}
|
||||
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(
|
||||
terms=actor_terms, concatenate_terms=True,
|
||||
enable_corruption=True, history_length=6,
|
||||
),
|
||||
"critic": ObservationGroupCfg(
|
||||
terms=critic_terms, concatenate_terms=True, enable_corruption=False,
|
||||
),
|
||||
}
|
||||
|
||||
# ------------------
|
||||
# Actions & Commands
|
||||
# ------------------
|
||||
actions: dict[str, ActionTermCfg] = {
|
||||
"leg_joint_pos": JointPositionDelayedLowPassActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"),
|
||||
scale={".*_hip_abduction_joint": 0.125, "^(?!.*_hip_abduction_joint).*": 0.25}, use_default_offset=True,
|
||||
control_frequency=50.0, cut_off_frequency=5.0,
|
||||
min_delay=0, max_delay=2,
|
||||
),
|
||||
"wheel_joint_vel": JointVelocityDelayedLowPassActionCfg(
|
||||
entity_name="robot", actuator_names=(".*_wheel_joint",),
|
||||
scale=5.0, offset=0.0, use_default_offset=False,
|
||||
control_frequency=50.0, cut_off_frequency=15.0,
|
||||
min_delay=0, max_delay=2,
|
||||
),
|
||||
}
|
||||
|
||||
commands: dict[str, CommandTermCfg] = {
|
||||
"twist": UniformThresholdVelocityCommandCfg(
|
||||
entity_name="robot", resampling_time_range=(10.0, 10.0),
|
||||
rel_standing_envs=0.15, rel_heading_envs=1.0, heading_command=True,
|
||||
heading_control_stiffness=0.6,
|
||||
rel_forward_envs=0.40,
|
||||
ranges=UniformThresholdVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-0.5, 0.5),
|
||||
ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
# ------------------
|
||||
# Domain Randomization (Events)
|
||||
# ------------------
|
||||
events = {
|
||||
"reset_scene": EventTermCfg(func=envs_mdp.reset_scene_to_default, mode="reset"),
|
||||
"reset_base": EventTermCfg(
|
||||
func=envs_mdp.reset_root_state_uniform, mode="reset",
|
||||
params={
|
||||
"pose_range": {"z": (0.30, 0.50), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.5, 0.5), "y": (-0.15, 0.15), "yaw": (-0.35, 0.35)},
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
),
|
||||
"push_robot": EventTermCfg(
|
||||
func=envs_mdp.push_by_setting_velocity, mode="interval",
|
||||
interval_range_s=(10.0, 15.0),
|
||||
params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}, "asset_cfg": SceneEntityCfg("robot")},
|
||||
),
|
||||
"base_com": EventTermCfg(
|
||||
func=envs_dr.body_com_offset, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
|
||||
"operation": "add", "ranges": {0: (-0.05, 0.05), 1: (-0.05, 0.05), 2: (-0.05, 0.05)}},
|
||||
),
|
||||
"encoder_bias": EventTermCfg(
|
||||
func=envs_dr.encoder_bias, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)), "bias_range": (-0.015, 0.015)},
|
||||
),
|
||||
"body_friction": EventTermCfg(
|
||||
func=envs_dr.geom_friction, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot", geom_names=(".*",)), "operation": "abs", "ranges": (0.3, 1.2)},
|
||||
),
|
||||
"actuator_stiffness": EventTermCfg(
|
||||
func=envs_dr.joint_stiffness, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.9, 1.1), "operation": "scale", "distribution": "log_uniform"},
|
||||
),
|
||||
"actuator_damping": EventTermCfg(
|
||||
func=envs_dr.joint_damping, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.9, 1.1), "operation": "scale", "distribution": "log_uniform"},
|
||||
),
|
||||
"actuator_effort_limit": EventTermCfg(
|
||||
func=envs_dr.effort_limits, mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
"effort_limit_range": (0.8, 1.0),
|
||||
"operation": "scale",
|
||||
"distribution": "uniform",
|
||||
},
|
||||
),
|
||||
"payload_mass": EventTermCfg(
|
||||
func=envs_dr.body_mass, mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
|
||||
"operation": "add",
|
||||
"ranges": (-1.0, 3.0),
|
||||
},
|
||||
),
|
||||
"continuous_disturbance": EventTermCfg(
|
||||
func=apply_continuous_disturbance, mode="step",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
|
||||
"force_range": (-15.0, 15.0),
|
||||
"torque_range": (-10.0, 10.0),
|
||||
"resample_time_range": (0.5, 2.0),
|
||||
"time_constant": 0.5,
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
# ------------------
|
||||
# Rewards Setup
|
||||
# ------------------
|
||||
rewards = {
|
||||
"track_lin_vel": RewardTermCfg(func=track_linear_velocity, weight=2.5, params={"std": 0.5, "command_name": "twist"}),
|
||||
"track_ang_vel": RewardTermCfg(func=track_angular_velocity, weight=2.5, params={"std": 0.5, "command_name": "twist"}),
|
||||
"upright": RewardTermCfg(func=velocity_mdp.upright, weight=1.0, params={"std": 0.5, "asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))}),
|
||||
"base_height_l2": RewardTermCfg(func=base_height_l2, weight=-2.0, params={"target_height": 0.36}),
|
||||
"body_ang_vel": RewardTermCfg(func=velocity_mdp.body_angular_velocity_penalty, weight=-0.1, params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))}),
|
||||
"is_terminated": RewardTermCfg(func=envs_mdp.is_terminated, weight=-200.0),
|
||||
"joint_torques": RewardTermCfg(func=envs_mdp.joint_torques_l2, weight=-2.0e-4),
|
||||
"joint_acc": RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-7),
|
||||
"action_rate": RewardTermCfg(func=envs_mdp.action_rate_l2, weight=-0.01),
|
||||
"joint_pos_limits": RewardTermCfg(func=envs_mdp.joint_pos_limits, weight=-10.0),
|
||||
"wheel_roll_tracking": RewardTermCfg(func=wheel_roll_tracking, weight=2.0, params={"command_name": "twist", "wheel_radius": 0.10, "wheel_track": 0.32, "std": 3.0, "asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))}),
|
||||
"wheel_contact_bonus": RewardTermCfg(func=contact_fraction_reward, weight=0.5, params={"sensor_name": "feet_ground_contact"}),
|
||||
"feet_air_time": RewardTermCfg(func=velocity_mdp.feet_air_time, weight=0.5, params={"sensor_name": "feet_ground_contact", "threshold_min": 0.1, "threshold_max": 0.5, "command_name": "twist", "command_threshold": 0.1}),
|
||||
"leg_motion_penalty": RewardTermCfg(func=adaptive_leg_motion_penalty, weight=-0.02, params={"command_name": "twist", "sensor_name": "feet_ground_contact", "command_threshold": 0.05, "tilt_relax_start": 0.08, "tilt_relax_end": 0.30, "contact_target": 0.85, "min_penalty_scale": 0.2, "asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))}),
|
||||
"stand_still": RewardTermCfg(func=stand_still, weight=-0.2, params={"command_name": "twist", "command_threshold": 0.1}),
|
||||
"body_collision": RewardTermCfg(func=velocity_mdp.self_collision_cost, weight=-1.0, params={"sensor_name": "body_collision"}),
|
||||
}
|
||||
|
||||
# ------------------
|
||||
# Terminations
|
||||
# ------------------
|
||||
terminations = {
|
||||
"time_out": TerminationTermCfg(func=envs_mdp.time_out, time_out=True),
|
||||
"bad_orientation": TerminationTermCfg(func=envs_mdp.bad_orientation, params={"limit_angle": 1.0}),
|
||||
"base_ground_contact": TerminationTermCfg(func=velocity_mdp.illegal_contact, params={"sensor_name": "base_ground_contact"}),
|
||||
"nan_detection": TerminationTermCfg(func=envs_mdp.nan_detection),
|
||||
}
|
||||
|
||||
curriculum: dict[str, CurriculumTermCfg] = {}
|
||||
|
||||
metrics = {"mean_leg_action_acc": MetricsTermCfg(func=velocity_mdp.mean_action_acc)}
|
||||
|
||||
return ManagerBasedRlEnvCfg(
|
||||
scene=SceneCfg(
|
||||
num_envs=2048, env_spacing=2.5,
|
||||
terrain=TerrainEntityCfg(terrain_type="generator", terrain_generator=TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0), border_width=20.0, num_rows=10, num_cols=20,
|
||||
sub_terrains={"flat": BoxFlatTerrainCfg(proportion=1.0)},
|
||||
)),
|
||||
sensors=(feet_ground_cfg, base_ground_cfg, body_collision_cfg, terrain_scan),
|
||||
),
|
||||
commands=commands, actions=actions, observations=observations,
|
||||
rewards=rewards, terminations=terminations, events=events,
|
||||
metrics=metrics, curriculum=curriculum, decimation=10, episode_length_s=20.0,
|
||||
sim=SimulationCfg(mujoco=MujocoCfg(impratio=100, cone="elliptic")),
|
||||
viewer=ViewerConfig(body_name="base_link", distance=3.0, elevation=-20.0, azimuth=45.0),
|
||||
)
|
||||
|
||||
|
||||
def flat_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
"""Flat ground training and evaluation configuration."""
|
||||
cfg = _make_base_env_cfg()
|
||||
cfg.scene.entities = {"robot": get_robot_cfg()}
|
||||
if play:
|
||||
cfg.episode_length_s = int(1e9)
|
||||
cfg.observations["actor"].enable_corruption = False
|
||||
cfg.events.pop("push_robot", None)
|
||||
cfg.curriculum = {}
|
||||
return cfg
|
||||
|
||||
|
||||
def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
"""Rough terrains configuration for general wheel-legged navigation."""
|
||||
enable_only_positive_rewards()
|
||||
|
||||
cfg = _make_base_env_cfg()
|
||||
cfg.scene.entities = {"robot": get_robot_cfg()}
|
||||
|
||||
# ------------------
|
||||
# Terrain Generator & Curriculum
|
||||
# ------------------
|
||||
cfg.scene.terrain = TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0), border_width=20.0, num_rows=10, num_cols=20, curriculum=True,
|
||||
sub_terrains={
|
||||
"flat": BoxFlatTerrainCfg(proportion=0.05, size=(8.0, 8.0)),
|
||||
"pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.25, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.10, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"random_grid": BoxRandomGridTerrainCfg(proportion=0.1, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
|
||||
"random_rough": HfRandomUniformTerrainCfg(proportion=0.05, noise_range=(0.0, 0.06), noise_step=0.01, horizontal_scale=0.20, downsampled_scale=0.20, border_width=0.25, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.05, height_range=(0.0, 0.06), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"rc_wall": RCWallTerrainCfg(proportion=0.25, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
|
||||
"sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.15, slope_range=(0.052, 0.325), platform_width=2.0, border_width=0.25, base_thickness_ratio=100.0, horizontal_scale=0.20, size=(8.0, 8.0)),
|
||||
},
|
||||
),
|
||||
max_init_terrain_level=0,
|
||||
)
|
||||
|
||||
# Disable default velocity stages command and bind strict velocity terrain curriculum
|
||||
cfg.curriculum.pop("command_vel", None)
|
||||
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=terrain_levels_vel_strict, params={"command_name": "twist"})
|
||||
|
||||
cfg.commands["twist"].heading_command = True
|
||||
cfg.commands["twist"].rel_heading_envs = 0.5
|
||||
cfg.commands["twist"].heading_control_stiffness = 0.6
|
||||
cfg.commands["twist"].ranges.heading = (-math.pi, math.pi)
|
||||
cfg.commands["twist"].rel_standing_envs = 0.2
|
||||
|
||||
# ------------------
|
||||
# Startup & Reset Randomizations
|
||||
# ------------------
|
||||
cfg.events["joint_friction"] = EventTermCfg(func=envs_dr.joint_friction, mode="startup", params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.7, 1.3), "operation": "scale"})
|
||||
cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.1), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
|
||||
|
||||
cfg.events["reset_base"] = EventTermCfg(
|
||||
func=envs_mdp.reset_root_state_uniform, mode="reset",
|
||||
params={
|
||||
"pose_range": {"z": (0.40, 0.45), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.2, 0.2), "y": (-0.1, 0.1), "yaw": (-0.2, 0.2)},
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
)
|
||||
cfg.events["push_robot"] = EventTermCfg(
|
||||
func=envs_mdp.push_by_setting_velocity, mode="interval",
|
||||
interval_range_s=(5.0, 10.0),
|
||||
params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}, "asset_cfg": SceneEntityCfg("robot")},
|
||||
)
|
||||
|
||||
# ------------------
|
||||
# Rewards Integration
|
||||
# ------------------
|
||||
cfg.rewards["track_lin_vel"] = RewardTermCfg(
|
||||
func=track_linear_velocity_l1,
|
||||
weight=4.5,
|
||||
params={"std": 0.5, "command_name": "twist"}
|
||||
)
|
||||
cfg.rewards["track_ang_vel"].weight = 2.0
|
||||
cfg.rewards["track_ang_vel"].params["std"] = 0.5
|
||||
|
||||
cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-0.5) # 🌟 增强垂直速度惩罚，抑制越障后的惯性暴冲
|
||||
cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=velocity_mdp.body_angular_velocity_penalty, weight=-0.3, params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))}) # 🌟 增强角速度惩罚，防止突发性翻转/后仰
|
||||
|
||||
cfg.rewards.pop("upright", None)
|
||||
cfg.rewards["roll_penalty"] = RewardTermCfg(
|
||||
func=upright_roll_only,
|
||||
weight=-1.0,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")}
|
||||
)
|
||||
|
||||
# 🌟 限制俯仰角死区（Pitch Dead-zone）：允许正常爬坡时有最大 29 度（0.50 rad）的仰角，但严厉惩罚超过该仰角的“前轮悬空暴冲/后翻”
|
||||
cfg.rewards["pitch_penalty"] = RewardTermCfg(
|
||||
func=pitch_control_penalty,
|
||||
weight=-1.5,
|
||||
params={"max_pitch_rad": 0.50, "asset_cfg": SceneEntityCfg("robot")}
|
||||
)
|
||||
|
||||
# 动态课程奖励与动作惩罚衰减
|
||||
cfg.rewards.pop("terrain_level_bonus", None)
|
||||
cfg.rewards.pop("action_rate", None)
|
||||
cfg.rewards["action_rate_curriculum"] = RewardTermCfg(func=action_rate_curriculum_l2, weight=-0.005)
|
||||
|
||||
cfg.rewards["joint_torques"].weight = -1e-4
|
||||
cfg.rewards.pop("joint_acc", None)
|
||||
cfg.rewards["leg_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-7, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))})
|
||||
cfg.rewards["wheel_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-9, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))})
|
||||
|
||||
|
||||
cfg.rewards["joint_pos_limits"].weight = -0.2
|
||||
cfg.rewards.pop("leg_motion_penalty", None)
|
||||
cfg.rewards["is_terminated"].weight = 0.0
|
||||
cfg.rewards.pop("leg_symmetry", None)
|
||||
cfg.rewards["joint_mirror"] = RewardTermCfg(
|
||||
func=joint_mirror,
|
||||
weight=-0.05,
|
||||
params={
|
||||
"mirror_joints": [
|
||||
["fl_(hip_abduction|hip_pitch|knee)_joint", "rr_(hip_abduction|hip_pitch|knee)_joint"],
|
||||
["fr_(hip_abduction|hip_pitch|knee)_joint", "rl_(hip_abduction|hip_pitch|knee)_joint"]
|
||||
],
|
||||
"asset_cfg": SceneEntityCfg("robot")
|
||||
}
|
||||
)
|
||||
|
||||
# 移除 variable_posture 及其产生的静止奖励陷阱，换用极轻微的偏离惩罚
|
||||
cfg.rewards.pop("stand_still", None)
|
||||
cfg.rewards["stand_still"] = RewardTermCfg(func=stand_still, weight=-2.0, params={"command_name": "twist", "command_threshold": 0.1})
|
||||
|
||||
cfg.rewards.pop("hip_deviation", None)
|
||||
cfg.rewards.pop("variable_posture", None)
|
||||
|
||||
cfg.rewards.pop("joint_deviation_l2", None)
|
||||
cfg.rewards["joint_pos_penalty"] = RewardTermCfg(
|
||||
func=joint_pos_penalty,
|
||||
weight=-0.8,
|
||||
params={
|
||||
"stand_still_scale": 5.0,
|
||||
"velocity_threshold": 0.5,
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint")),
|
||||
"command_name": "twist"
|
||||
}
|
||||
)
|
||||
|
||||
cfg.rewards["feet_contact_without_cmd"] = RewardTermCfg(
|
||||
func=feet_contact_without_cmd,
|
||||
weight=0.1,
|
||||
params={"command_name": "twist", "sensor_name": "feet_ground_contact"}
|
||||
)
|
||||
|
||||
cfg.rewards["base_height_l2"].weight = -0.5
|
||||
cfg.rewards["base_height_l2"].params["sensor_cfg"] = SceneEntityCfg("height_scanner")
|
||||
|
||||
# 恢复机身碰撞惩罚为-1.0，逼迫机器人高抬腿跨越障碍，防止拖地
|
||||
cfg.rewards.pop("body_collision", None)
|
||||
cfg.rewards["body_collision"] = RewardTermCfg(func=velocity_mdp.self_collision_cost, weight=-1.0, params={"sensor_name": "body_collision"})
|
||||
|
||||
# 🌟 严厉惩罚机身/胸部碰撞（防止硬撞高墙），逼迫机器人学会用前轮触墙并主动抬腿攀爬的“触觉反射”
|
||||
cfg.rewards["base_collision"] = RewardTermCfg(
|
||||
func=velocity_mdp.self_collision_cost,
|
||||
weight=-5.0,
|
||||
params={"sensor_name": "base_ground_contact"}
|
||||
)
|
||||
|
||||
# 彻底移除机身俯仰约束，允许机器人抬头爬高? cfg.rewards.pop("flat_orientation", None)
|
||||
|
||||
cfg.rewards.pop("feet_air_time", None)
|
||||
|
||||
# Remove non-applicable rewards
|
||||
for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel"):
|
||||
cfg.rewards.pop(key, None)
|
||||
|
||||
cfg.episode_length_s = 30.0
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
|
||||
# 移除 orientation 终止，允许机器人翻倒以学习回复
|
||||
cfg.terminations.pop("bad_orientation", None)
|
||||
# 移除 base_ground_contact 终止，越障时机身会碰到障碍物
|
||||
cfg.terminations.pop("base_ground_contact", None)
|
||||
|
||||
cfg.seed = 42
|
||||
if cfg.scene.terrain is not None:
|
||||
cfg.scene.terrain.num_envs = 2048
|
||||
cfg.scene.terrain.env_spacing = 2.5
|
||||
|
||||
if play:
|
||||
cfg.episode_length_s = int(1e9)
|
||||
cfg.observations["actor"].enable_corruption = False
|
||||
cfg.events.pop("push_robot", None)
|
||||
cfg.curriculum = {}
|
||||
if cfg.scene.terrain is not None and cfg.scene.terrain.terrain_generator is not None:
|
||||
cfg.scene.terrain.terrain_generator.curriculum = False
|
||||
cfg.scene.terrain.terrain_generator.num_cols = 5
|
||||
cfg.scene.terrain.terrain_generator.num_rows = 5
|
||||
cfg.scene.terrain.terrain_generator.border_width = 10.0
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
def crawl_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
"""Crawling and ducking configuration for climbing obstacles under low clearing heights."""
|
||||
enable_only_positive_rewards()
|
||||
|
||||
cfg = _make_base_env_cfg()
|
||||
cfg.scene.entities = {"robot": get_robot_crawl_cfg()}
|
||||
|
||||
# ------------------
|
||||
# Terrain Definition
|
||||
# ------------------
|
||||
cfg.scene.terrain = TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0), border_width=20.0, num_rows=10, num_cols=20, curriculum=True,
|
||||
sub_terrains={
|
||||
"flat": BoxFlatTerrainCfg(proportion=0.25),
|
||||
"rc_low_bar": RCLowBarTerrainCfg(proportion=0.35, clearance_range=(0.24, 0.32)),
|
||||
"random_grid": BoxRandomGridTerrainCfg(proportion=0.20, grid_width=0.45, grid_height_range=(0.0, 0.05)),
|
||||
"perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.20, height_range=(0.0, 0.05), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0),
|
||||
},
|
||||
),
|
||||
max_init_terrain_level=0,
|
||||
)
|
||||
|
||||
# Disable command vel curriculum and setup base command ranges
|
||||
cfg.curriculum.pop("command_vel", None)
|
||||
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=velocity_mdp.terrain_levels_vel, params={"command_name": "twist"})
|
||||
|
||||
cfg.commands["twist"].heading_command = False
|
||||
cfg.commands["twist"].rel_heading_envs = 0.0
|
||||
cfg.commands["twist"].ranges.heading = None
|
||||
cfg.commands["twist"].rel_standing_envs = 0.05 # 减少静止比例（匍匐需持续运动�? cfg.commands["twist"].rel_forward_envs = 0.40 # 40% 纯前向（低杆地形全是直穿�?
|
||||
# ------------------
|
||||
# Events & Reset
|
||||
# ------------------
|
||||
cfg.events["joint_friction"] = EventTermCfg(func=envs_dr.joint_friction, mode="startup", params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.7, 1.3), "operation": "scale"})
|
||||
cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.1), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
|
||||
|
||||
cfg.events["reset_base"] = EventTermCfg(
|
||||
func=envs_mdp.reset_root_state_uniform, mode="reset",
|
||||
params={
|
||||
"pose_range": {"z": (0.00, 0.04), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.1, 0.1), "y": (-0.05, 0.05), "yaw": (-0.1, 0.1)},
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
)
|
||||
cfg.events["push_robot"] = EventTermCfg(
|
||||
func=envs_mdp.push_by_setting_velocity, mode="interval",
|
||||
interval_range_s=(5.0, 10.0),
|
||||
params={"velocity_range": {"x": (-0.3, 0.3), "y": (-0.3, 0.3)}, "asset_cfg": SceneEntityCfg("robot")},
|
||||
)
|
||||
|
||||
# ------------------
|
||||
# Rewards Integration
|
||||
# ------------------
|
||||
cfg.rewards["track_lin_vel"].weight = 3.0
|
||||
cfg.rewards["track_lin_vel"].params["std"] = 0.5
|
||||
cfg.rewards["track_ang_vel"].weight = 1.5
|
||||
cfg.rewards["track_ang_vel"].params["std"] = 0.5
|
||||
|
||||
cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-1.0)
|
||||
cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=velocity_mdp.body_angular_velocity_penalty, weight=-0.05, params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))})
|
||||
|
||||
cfg.rewards["upright"].weight = 1.0
|
||||
cfg.rewards["upright"].params["std"] = 0.5
|
||||
cfg.rewards["action_rate"].weight = -0.001
|
||||
cfg.rewards["joint_torques"].weight = -1e-4
|
||||
cfg.rewards["joint_acc"].weight = 0.0
|
||||
cfg.rewards["joint_pos_limits"].weight = -1.0
|
||||
cfg.rewards["leg_motion_penalty"].weight = -5.0
|
||||
cfg.rewards["is_terminated"].weight = -50.0
|
||||
|
||||
# Under-crawling height reward: maximum bonus when body stays under 0.22m
|
||||
cfg.rewards.pop("base_height_l2", None)
|
||||
cfg.rewards["crawl_height_reward"] = RewardTermCfg(
|
||||
func=crawl_height_reward,
|
||||
weight=1.5,
|
||||
params={"target_height": 0.22, "std": 0.05}
|
||||
)
|
||||
|
||||
cfg.rewards.pop("stand_still", None)
|
||||
|
||||
cfg.rewards.pop("hip_deviation", None)
|
||||
cfg.rewards["leg_joint_deviation"] = RewardTermCfg(
|
||||
func=joint_deviation_l2,
|
||||
weight=-15.0,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))},
|
||||
)
|
||||
|
||||
# 🌟 强力引入轮子滚动跟踪奖励，引导机器人完全依靠轮子的正向滚动进行平地/低矮处的推进
|
||||
cfg.rewards["wheel_roll_tracking"] = RewardTermCfg(
|
||||
func=wheel_roll_tracking,
|
||||
weight=4.0,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"wheel_radius": 0.10,
|
||||
"wheel_track": 0.32,
|
||||
"std": 3.0,
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))
|
||||
}
|
||||
)
|
||||
|
||||
if "body_collision" in cfg.rewards:
|
||||
cfg.rewards["body_collision"].weight = -0.1
|
||||
|
||||
cfg.rewards["flat_orientation"] = RewardTermCfg(func=flat_orientation_l2, weight=-1.0)
|
||||
|
||||
for key in ("wheel_roll_tracking", "feet_air_time", "wheel_contact_bonus", "body_ang_vel"):
|
||||
cfg.rewards.pop(key, None)
|
||||
|
||||
cfg.episode_length_s = 30.0
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
|
||||
# Loosen orientation bad threshold to 80 degrees for steep crawling tilts
|
||||
cfg.terminations["bad_orientation"].params["limit_angle"] = math.radians(80.0)
|
||||
|
||||
# Remove base ground contact termination to facilitate crawl under bars
|
||||
cfg.terminations.pop("base_ground_contact", None)
|
||||
|
||||
if play:
|
||||
cfg.episode_length_s = int(1e9)
|
||||
cfg.observations["actor"].enable_corruption = False
|
||||
cfg.events.pop("push_robot", None)
|
||||
cfg.curriculum = {}
|
||||
if cfg.scene.terrain is not None and cfg.scene.terrain.terrain_generator is not None:
|
||||
cfg.scene.terrain.terrain_generator.curriculum = False
|
||||
cfg.scene.terrain.terrain_generator.num_cols = 5
|
||||
cfg.scene.terrain.terrain_generator.num_rows = 5
|
||||
cfg.scene.terrain.terrain_generator.border_width = 10.0
|
||||
|
||||
return cfg
|
||||
@@ -0,0 +1,59 @@
|
||||
from mjlab.rl import (
|
||||
RslRlModelCfg,
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
|
||||
def flat_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
cfg = rough_ppo_runner_cfg()
|
||||
cfg.experiment_name = "robot_flat"
|
||||
cfg.max_iterations = 10_000
|
||||
return cfg
|
||||
|
||||
|
||||
def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
return RslRlOnPolicyRunnerCfg(
|
||||
actor=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=False,
|
||||
distribution_cfg={
|
||||
"class_name": "GaussianDistribution",
|
||||
"init_std": 1.0,
|
||||
"std_type": "log",
|
||||
},
|
||||
),
|
||||
critic=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=False,
|
||||
),
|
||||
algorithm=RslRlPpoAlgorithmCfg(
|
||||
value_loss_coef=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.002,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=2.0e-4,
|
||||
schedule="adaptive",
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
),
|
||||
experiment_name="robot_rough",
|
||||
save_interval=50,
|
||||
num_steps_per_env=24,
|
||||
max_iterations=15_000,
|
||||
)
|
||||
|
||||
|
||||
def crawl_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
cfg = rough_ppo_runner_cfg()
|
||||
cfg.experiment_name = "robot_crawl"
|
||||
return cfg
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
|
||||
|
||||
|
||||
class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
"""带死区过滤、高速侧向解耦以及地形自适应速度约束的速度命令生成器。
|
||||
|
||||
在楼梯和垂直短墙等障碍地形上，自动将 y 轴（横移）和 z 轴（偏航）控制约束置 0.0，
|
||||
强制机器人心无旁骛笔直冲锋越障，有效杜绝打滑和偏航翻滚跌落；而在平地、斜坡等普通地形上
|
||||
放开多向混合采样，训练机动转向能力。
|
||||
"""
|
||||
cfg: UniformThresholdVelocityCommandCfg
|
||||
|
||||
def __init__(self, cfg: UniformThresholdVelocityCommandCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
# 缓存地形类型的索引（台阶、反向台阶、垂直短墙），实现高容错动态查找
|
||||
self._climbing_indices = []
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
if terrain is not None and getattr(terrain.cfg, "terrain_generator", None) is not None:
|
||||
sub_terrain_names = list(terrain.cfg.terrain_generator.sub_terrains.keys())
|
||||
for name in ["pyramid_stairs", "pyramid_stairs_inv", "rc_wall"]:
|
||||
if name in sub_terrain_names:
|
||||
self._climbing_indices.append(sub_terrain_names.index(name))
|
||||
|
||||
def _resample_command(self, env_ids: torch.Tensor) -> None:
|
||||
# 1. 调用基类的标准采样
|
||||
super()._resample_command(env_ids)
|
||||
|
||||
# 2. 基础死区过滤：低于 0.2 m/s 时强制归零，区分静止和运动
|
||||
cmd_xy_norm = torch.norm(self.vel_command_b[env_ids, :2], dim=1)
|
||||
small_cmd_mask = cmd_xy_norm < 0.2
|
||||
small_cmd_ids = env_ids[small_cmd_mask]
|
||||
if len(small_cmd_ids) > 0:
|
||||
self.vel_command_b[small_cmd_ids, :] = 0.0
|
||||
self.vel_command_w[small_cmd_ids, :] = 0.0
|
||||
|
||||
# 3. 地形自适应重采样限制：若在爬行地形，强制纯前进方向且速度 >= 0.3 m/s
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
terrain_types = getattr(terrain, "terrain_types", None)
|
||||
if terrain_types is not None and len(self._climbing_indices) > 0:
|
||||
is_climbing = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
for idx in self._climbing_indices:
|
||||
is_climbing |= (terrain_types == idx)
|
||||
|
||||
climbing_env_ids = env_ids[is_climbing[env_ids]]
|
||||
if len(climbing_env_ids) > 0:
|
||||
# 强制 x 方向为正（向前越障），y 和 z 轴指令清零
|
||||
self.vel_command_b[climbing_env_ids, 0] = self.vel_command_b[climbing_env_ids, 0].abs().clamp(min=0.3)
|
||||
self.vel_command_b[climbing_env_ids, 1] = 0.0
|
||||
self.vel_command_b[climbing_env_ids, 2] = 0.0
|
||||
self.is_heading_env[climbing_env_ids] = True
|
||||
|
||||
def _update_command(self) -> None:
|
||||
# 调用基类的每步更新
|
||||
super()._update_command()
|
||||
|
||||
# 4. 每步更新时强力约束：若处于台阶/短墙爬高地形，强制 y 轴横移持续为 0.0，允许温和的偏航纠偏对齐台阶
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
terrain_types = getattr(terrain, "terrain_types", None)
|
||||
if terrain_types is not None and len(self._climbing_indices) > 0:
|
||||
is_climbing = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
for idx in self._climbing_indices:
|
||||
is_climbing |= (terrain_types == idx)
|
||||
|
||||
climbing_env_ids = is_climbing.nonzero(as_tuple=False).flatten()
|
||||
if len(climbing_env_ids) > 0:
|
||||
self.vel_command_b[climbing_env_ids, 1] = 0.0
|
||||
# 🌟 允许微弱的偏航纠偏，将偏航指令限制在温和的 [-0.3, 0.3] 区间，防止过度甩尾，但保证能修正航向
|
||||
self.vel_command_b[climbing_env_ids, 2] = torch.clip(
|
||||
self.cfg.heading_control_stiffness * self.heading_error[climbing_env_ids],
|
||||
min=-0.3,
|
||||
max=0.3
|
||||
)
|
||||
|
||||
# 5. 高速侧向解耦（适用于平地/斜坡等混合路面）：当前进速度 >= 0.8 m/s 时，清空侧向指令，防止高速甩尾甩飞
|
||||
high_speed_mask = self.vel_command_b[:, 0].abs() >= 0.8
|
||||
high_speed_ids = high_speed_mask.nonzero(as_tuple=False).flatten()
|
||||
if len(high_speed_ids) > 0:
|
||||
self.vel_command_b[high_speed_ids, 1] = 0.0
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class UniformThresholdVelocityCommandCfg(UniformVelocityCommandCfg):
|
||||
class_type: type = UniformThresholdVelocityCommand
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Adaptive command curriculum based on tracking reward performance."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, cast
|
||||
import torch
|
||||
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
class adaptive_command_vel:
|
||||
"""Expand velocity command ranges dynamically when tracking performance exceeds threshold.
|
||||
|
||||
Uses the reward manager's per-step reward buffer directly to compute the mean raw
|
||||
tracking reward (in [0, 1] for exponential-based rewards).
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRlEnv):
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
|
||||
p = cfg.params
|
||||
self._command_name: str = p["command_name"]
|
||||
self._reward_name: str = p["reward_name"]
|
||||
self._upper_threshold: float = p.get("upper_threshold", 0.8)
|
||||
self._lower_threshold: float = p.get("lower_threshold", 0.4)
|
||||
self._delta: float = p.get("delta", 0.2)
|
||||
self._max_lin_vel_x: tuple = tuple(p.get("max_lin_vel_x", (-1.0, 1.0)))
|
||||
self._max_lin_vel_y: tuple = tuple(p.get("max_lin_vel_y", (-0.5, 0.5)))
|
||||
self._max_ang_vel_z: tuple = tuple(p.get("max_ang_vel_z", (-1.0, 1.0)))
|
||||
self._min_range: float = 0.3
|
||||
|
||||
command_term = env.command_manager.get_term(self._command_name)
|
||||
self._cfg = cast(UniformVelocityCommandCfg, command_term.cfg)
|
||||
|
||||
# Locate the index of the target reward term
|
||||
self._reward_idx = list(env.reward_manager._term_names).index(self._reward_name)
|
||||
self._reward_weight = env.reward_manager.get_term_cfg(self._reward_name).weight
|
||||
|
||||
# Exponential moving average (EMA) smoothing to prevent oscillations
|
||||
self._ema = 0.5
|
||||
self._running_mean = 0.0
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
**kwargs,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
# Retrieve raw step rewards from step_reward buffer
|
||||
# step_reward[:, idx] = raw * weight (already scaled by dt)
|
||||
step_rewards = env.reward_manager._step_reward[:, self._reward_idx]
|
||||
# raw = step_reward / weight -> [0, 1] for exponential rewards
|
||||
mean_raw = (torch.mean(step_rewards) / self._reward_weight).item()
|
||||
|
||||
# Apply EMA smoothing to prevent jitter
|
||||
self._running_mean = self._ema * mean_raw + (1 - self._ema) * self._running_mean
|
||||
|
||||
if self._running_mean > self._upper_threshold:
|
||||
self._expand()
|
||||
elif self._running_mean < self._lower_threshold:
|
||||
self._shrink()
|
||||
|
||||
return self._log()
|
||||
|
||||
def _expand(self):
|
||||
lo, hi = self._cfg.ranges.lin_vel_x
|
||||
self._cfg.ranges.lin_vel_x = (
|
||||
max(lo - self._delta, self._max_lin_vel_x[0]),
|
||||
min(hi + self._delta, self._max_lin_vel_x[1]),
|
||||
)
|
||||
lo, hi = self._cfg.ranges.lin_vel_y
|
||||
self._cfg.ranges.lin_vel_y = (
|
||||
max(lo - self._delta * 0.5, self._max_lin_vel_y[0]),
|
||||
min(hi + self._delta * 0.5, self._max_lin_vel_y[1]),
|
||||
)
|
||||
lo, hi = self._cfg.ranges.ang_vel_z
|
||||
self._cfg.ranges.ang_vel_z = (
|
||||
max(lo - self._delta, self._max_ang_vel_z[0]),
|
||||
min(hi + self._delta, self._max_ang_vel_z[1]),
|
||||
)
|
||||
|
||||
def _shrink(self):
|
||||
lo, hi = self._cfg.ranges.lin_vel_x
|
||||
new_lo = min(lo + self._delta, -self._min_range)
|
||||
new_hi = max(hi - self._delta, self._min_range)
|
||||
if new_hi - new_lo >= 2 * self._min_range:
|
||||
self._cfg.ranges.lin_vel_x = (new_lo, new_hi)
|
||||
|
||||
def _log(self) -> dict[str, torch.Tensor]:
|
||||
return {
|
||||
"lin_vel_x_min": torch.tensor(self._cfg.ranges.lin_vel_x[0]),
|
||||
"lin_vel_x_max": torch.tensor(self._cfg.ranges.lin_vel_x[1]),
|
||||
"ang_vel_z_max": torch.tensor(self._cfg.ranges.ang_vel_z[1]),
|
||||
}
|
||||
|
||||
|
||||
def terrain_levels_vel_strict(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Velocity-based terrain curriculum aligned with standard legged_gym logic.
|
||||
|
||||
Upgrade: robot travels beyond half the terrain tile width (> 4m).
|
||||
Downgrade: actual distance < 50% of commanded target distance.
|
||||
|
||||
This matches DreamWaQ / HIMLoco / LocoLeggedWheel curriculum behaviour:
|
||||
- Promotion is easy (any traversal past 4m qualifies).
|
||||
- Demotion requires consistently failing to cover half the expected distance.
|
||||
"""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
terrain = env.scene.terrain
|
||||
assert terrain is not None
|
||||
terrain_generator = terrain.cfg.terrain_generator
|
||||
assert terrain_generator is not None
|
||||
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None
|
||||
|
||||
# Horizontal displacement from episode start
|
||||
distance = torch.norm(
|
||||
asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
cmd_speed = torch.norm(command[env_ids, :2], dim=1)
|
||||
|
||||
# Upgrade: crossed half the tile width
|
||||
move_up = distance > terrain_generator.size[0] / 2
|
||||
|
||||
# Downgrade: traveled less than 33% of commanded target distance.
|
||||
# Absolute threshold ≈ cmd_speed × 10m, identical to DreamWaQ / HIMLoco / LocoLeggedWheel
|
||||
# which use 50% × 20s episode = 10m. Adjusted for the longer 30s episode here.
|
||||
move_down = (distance < cmd_speed * env.max_episode_length_s * 0.33) & ~move_up
|
||||
|
||||
|
||||
|
||||
terrain.update_env_origins(env_ids, move_up, move_down)
|
||||
|
||||
levels = terrain.terrain_levels.float()
|
||||
result: dict[str, torch.Tensor] = {
|
||||
"mean": torch.mean(levels),
|
||||
"max": torch.max(levels),
|
||||
}
|
||||
|
||||
sub_terrain_names = list(terrain_generator.sub_terrains.keys())
|
||||
terrain_origins = terrain.terrain_origins
|
||||
assert terrain_origins is not None
|
||||
num_cols = terrain_origins.shape[1]
|
||||
if num_cols == len(sub_terrain_names):
|
||||
types = terrain.terrain_types
|
||||
for i, name in enumerate(sub_terrain_names):
|
||||
mask = types == i
|
||||
if mask.any():
|
||||
result[name] = torch.mean(levels[mask])
|
||||
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Continuous low-frequency external wrench disturbance for mjlab."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.utils.lab_api.math import sample_uniform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
class apply_continuous_disturbance:
|
||||
"""Apply smooth, continuous low-frequency force and torque disturbances to bodies.
|
||||
|
||||
Achieved by sampling random target forces and torques periodically, and interpolating
|
||||
towards them via a 1st-order low-pass filter (exponential smoothing).
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, env: ManagerBasedRlEnv):
|
||||
self._asset: Entity = env.scene[cfg.params["asset_cfg"].name]
|
||||
self._body_ids = cfg.params["asset_cfg"].body_ids
|
||||
self._num_envs = env.num_envs
|
||||
self._device = env.device
|
||||
self._step_dt = env.step_dt
|
||||
|
||||
self._num_bodies = (
|
||||
len(self._body_ids)
|
||||
if isinstance(self._body_ids, list)
|
||||
else self._asset.num_bodies
|
||||
)
|
||||
|
||||
# Disturbance states: shape (num_envs, num_bodies, 3)
|
||||
self._current_force = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device)
|
||||
self._current_torque = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device)
|
||||
self._target_force = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device)
|
||||
self._target_torque = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device)
|
||||
|
||||
# Periodic resampling timers: shape (num_envs,)
|
||||
self._timer = torch.zeros(self._num_envs, device=self._device)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor | None,
|
||||
force_range: tuple[float, float],
|
||||
torque_range: tuple[float, float],
|
||||
resample_time_range: tuple[float, float],
|
||||
time_constant: float,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
) -> None:
|
||||
"""Tick disturbance state: interpolate towards targets, periodically resample targets.
|
||||
|
||||
Use with mode="step".
|
||||
"""
|
||||
del env, env_ids, asset_cfg # Unused, step events act on all envs.
|
||||
dt = self._step_dt
|
||||
|
||||
# Decrement timers
|
||||
self._timer -= dt
|
||||
|
||||
# Identify envs that need resampling
|
||||
resample = self._timer <= 0
|
||||
if resample.any():
|
||||
resample_ids = resample.nonzero(as_tuple=False).squeeze(-1)
|
||||
n = len(resample_ids)
|
||||
|
||||
# Sample new targets
|
||||
size = (n, self._num_bodies, 3)
|
||||
self._target_force[resample_ids] = sample_uniform(*force_range, size, self._device)
|
||||
self._target_torque[resample_ids] = sample_uniform(*torque_range, size, self._device)
|
||||
|
||||
# Sample new timer durations
|
||||
t_low, t_high = resample_time_range
|
||||
self._timer[resample_ids] = (
|
||||
torch.rand(n, device=self._device) * (t_high - t_low) + t_low
|
||||
)
|
||||
|
||||
# Exponential smoothing step: x_new = (1 - alpha) * x + alpha * x_target
|
||||
alpha = 1.0 - math.exp(-dt / time_constant)
|
||||
|
||||
self._current_force = (1.0 - alpha) * self._current_force + alpha * self._target_force
|
||||
self._current_torque = (1.0 - alpha) * self._current_torque + alpha * self._target_torque
|
||||
|
||||
# Apply smooth wrenches to simulation
|
||||
all_env_ids = torch.arange(self._num_envs, device=self._device)
|
||||
self._asset.write_external_wrench_to_sim(
|
||||
self._current_force, self._current_torque, env_ids=all_env_ids, body_ids=self._body_ids
|
||||
)
|
||||
|
||||
def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
|
||||
if env_ids is None:
|
||||
env_ids = slice(None)
|
||||
|
||||
# Reset states
|
||||
self._current_force[env_ids] = 0.0
|
||||
self._current_torque[env_ids] = 0.0
|
||||
self._target_force[env_ids] = 0.0
|
||||
self._target_torque[env_ids] = 0.0
|
||||
self._timer[env_ids] = 0.0
|
||||
|
||||
if isinstance(env_ids, slice):
|
||||
reset_ids = torch.arange(self._num_envs, device=self._device)[env_ids]
|
||||
else:
|
||||
reset_ids = env_ids
|
||||
|
||||
if len(reset_ids) > 0:
|
||||
zeros = torch.zeros((len(reset_ids), self._num_bodies, 3), device=self._device)
|
||||
self._asset.write_external_wrench_to_sim(
|
||||
zeros, zeros, env_ids=reset_ids, body_ids=self._body_ids
|
||||
)
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Low-pass filtered action terms for mjlab (native implementation).
|
||||
|
||||
Provides FIR low-pass filtering on the raw policy output before applying
|
||||
scale/offset. This smooths control signals and reduces sim-to-real gap.
|
||||
|
||||
- Leg joints: 5 Hz cutoff (slow, smooth)
|
||||
- Wheel joints: 15 Hz cutoff (faster response needed)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.envs.mdp.actions.actions import (
|
||||
JointPositionAction,
|
||||
JointPositionActionCfg,
|
||||
JointVelocityAction,
|
||||
JointVelocityActionCfg,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
def _lowpass_weights(control_freq: float, cutoff_freq: float) -> list[float]:
|
||||
"""Compute 1st-order IIR low-pass filter weights.
|
||||
|
||||
alpha = 1 - exp(-2π * f_c / f_s)
|
||||
filtered = alpha * current + (1 - alpha) * previous
|
||||
"""
|
||||
alpha = 1.0 - math.exp(-2.0 * math.pi * cutoff_freq / control_freq)
|
||||
return [alpha, 1.0 - alpha]
|
||||
|
||||
|
||||
# -- Position low-pass --
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class JointPositionLowPassActionCfg(JointPositionActionCfg):
|
||||
"""Joint position action with 1st-order low-pass filter on policy output."""
|
||||
|
||||
control_frequency: float = 50.0
|
||||
"""Control frequency in Hz (= 1 / (dt * decimation))."""
|
||||
cut_off_frequency: float = 5.0
|
||||
"""Cut-off frequency in Hz for the low-pass filter."""
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> JointPositionLowPassAction:
|
||||
return JointPositionLowPassAction(self, env)
|
||||
|
||||
|
||||
class JointPositionLowPassAction(JointPositionAction):
|
||||
"""Applies a 1st-order low-pass filter to raw actions before processing."""
|
||||
|
||||
def __init__(self, cfg: JointPositionLowPassActionCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
self._weights = _lowpass_weights(cfg.control_frequency, cfg.cut_off_frequency)
|
||||
self._prev_raw = torch.zeros_like(self._raw_actions)
|
||||
|
||||
def process_actions(self, actions: torch.Tensor):
|
||||
filtered = self._weights[0] * actions + self._weights[1] * self._prev_raw
|
||||
self._prev_raw[:] = actions
|
||||
super().process_actions(filtered)
|
||||
|
||||
def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
|
||||
self._prev_raw[env_ids] = 0.0
|
||||
super().reset(env_ids)
|
||||
|
||||
|
||||
# -- Velocity low-pass --
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class JointVelocityLowPassActionCfg(JointVelocityActionCfg):
|
||||
"""Joint velocity action with 1st-order low-pass filter on policy output."""
|
||||
|
||||
control_frequency: float = 50.0
|
||||
"""Control frequency in Hz (= 1 / (dt * decimation))."""
|
||||
cut_off_frequency: float = 15.0
|
||||
"""Cut-off frequency in Hz for the low-pass filter."""
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> JointVelocityLowPassAction:
|
||||
return JointVelocityLowPassAction(self, env)
|
||||
|
||||
|
||||
class JointVelocityLowPassAction(JointVelocityAction):
|
||||
"""Applies a 1st-order low-pass filter to raw actions before processing."""
|
||||
|
||||
def __init__(self, cfg: JointVelocityLowPassActionCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
self._weights = _lowpass_weights(cfg.control_frequency, cfg.cut_off_frequency)
|
||||
self._prev_raw = torch.zeros_like(self._raw_actions)
|
||||
|
||||
def process_actions(self, actions: torch.Tensor):
|
||||
filtered = self._weights[0] * actions + self._weights[1] * self._prev_raw
|
||||
self._prev_raw[:] = actions
|
||||
super().process_actions(filtered)
|
||||
|
||||
def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
|
||||
self._prev_raw[env_ids] = 0.0
|
||||
super().reset(env_ids)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class JointPositionDelayedLowPassActionCfg(JointPositionLowPassActionCfg):
|
||||
"""Joint position action with random delay and 1st-order low-pass filter."""
|
||||
|
||||
min_delay: int = 0
|
||||
max_delay: int = 2
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> JointPositionDelayedLowPassAction:
|
||||
return JointPositionDelayedLowPassAction(self, env)
|
||||
|
||||
|
||||
class JointPositionDelayedLowPassAction(JointPositionLowPassAction):
|
||||
"""Applies a random delay and a 1st-order low-pass filter to raw actions."""
|
||||
|
||||
def __init__(self, cfg: JointPositionDelayedLowPassActionCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
self.min_delay = cfg.min_delay
|
||||
self.max_delay = cfg.max_delay
|
||||
|
||||
# Buffer shape: (max_delay + 1, num_envs, action_dim)
|
||||
self._action_buffer = torch.zeros(
|
||||
self.max_delay + 1, self.num_envs, self.action_dim, device=self.device
|
||||
)
|
||||
|
||||
# Active delays for each env, shape: (num_envs,)
|
||||
self._active_delays = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
|
||||
self.reset_delays(torch.arange(self.num_envs, device=self.device))
|
||||
|
||||
def reset_delays(self, env_ids: torch.Tensor):
|
||||
if self.min_delay == self.max_delay:
|
||||
self._active_delays[env_ids] = self.min_delay
|
||||
else:
|
||||
sampled = torch.randint(
|
||||
self.min_delay,
|
||||
self.max_delay + 1,
|
||||
(len(env_ids),),
|
||||
dtype=torch.long,
|
||||
device=self.device,
|
||||
)
|
||||
self._active_delays[env_ids] = sampled
|
||||
|
||||
def process_actions(self, actions: torch.Tensor):
|
||||
if self.max_delay > 0:
|
||||
self._action_buffer = torch.cat(
|
||||
[self._action_buffer[1:], actions.unsqueeze(0)], dim=0
|
||||
)
|
||||
else:
|
||||
self._action_buffer[0] = actions
|
||||
|
||||
indices = self.max_delay - self._active_delays
|
||||
actions_delayed = self._action_buffer[
|
||||
indices, torch.arange(self.num_envs, device=self.device)
|
||||
]
|
||||
|
||||
filtered = self._weights[0] * actions_delayed + self._weights[1] * self._prev_raw
|
||||
self._prev_raw[:] = actions_delayed
|
||||
super(JointPositionLowPassAction, self).process_actions(filtered)
|
||||
|
||||
def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(self.num_envs, device=self.device)
|
||||
elif isinstance(env_ids, slice):
|
||||
env_ids = torch.arange(self.num_envs, device=self.device)[env_ids]
|
||||
|
||||
self.reset_delays(env_ids)
|
||||
self._action_buffer[:, env_ids] = 0.0
|
||||
self._prev_raw[env_ids] = 0.0
|
||||
super().reset(env_ids)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class JointVelocityDelayedLowPassActionCfg(JointVelocityLowPassActionCfg):
|
||||
"""Joint velocity action with random delay and 1st-order low-pass filter."""
|
||||
|
||||
min_delay: int = 0
|
||||
max_delay: int = 2
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> JointVelocityDelayedLowPassAction:
|
||||
return JointVelocityDelayedLowPassAction(self, env)
|
||||
|
||||
|
||||
class JointVelocityDelayedLowPassAction(JointVelocityLowPassAction):
|
||||
"""Applies a random delay and a 1st-order low-pass filter to raw actions."""
|
||||
|
||||
def __init__(self, cfg: JointVelocityDelayedLowPassActionCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
self.min_delay = cfg.min_delay
|
||||
self.max_delay = cfg.max_delay
|
||||
|
||||
# Buffer shape: (max_delay + 1, num_envs, action_dim)
|
||||
self._action_buffer = torch.zeros(
|
||||
self.max_delay + 1, self.num_envs, self.action_dim, device=self.device
|
||||
)
|
||||
|
||||
# Active delays for each env, shape: (num_envs,)
|
||||
self._active_delays = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
|
||||
self.reset_delays(torch.arange(self.num_envs, device=self.device))
|
||||
|
||||
def reset_delays(self, env_ids: torch.Tensor):
|
||||
if self.min_delay == self.max_delay:
|
||||
self._active_delays[env_ids] = self.min_delay
|
||||
else:
|
||||
sampled = torch.randint(
|
||||
self.min_delay,
|
||||
self.max_delay + 1,
|
||||
(len(env_ids),),
|
||||
dtype=torch.long,
|
||||
device=self.device,
|
||||
)
|
||||
self._active_delays[env_ids] = sampled
|
||||
|
||||
def process_actions(self, actions: torch.Tensor):
|
||||
if self.max_delay > 0:
|
||||
self._action_buffer = torch.cat(
|
||||
[self._action_buffer[1:], actions.unsqueeze(0)], dim=0
|
||||
)
|
||||
else:
|
||||
self._action_buffer[0] = actions
|
||||
|
||||
indices = self.max_delay - self._active_delays
|
||||
actions_delayed = self._action_buffer[
|
||||
indices, torch.arange(self.num_envs, device=self.device)
|
||||
]
|
||||
|
||||
filtered = self._weights[0] * actions_delayed + self._weights[1] * self._prev_raw
|
||||
self._prev_raw[:] = actions_delayed
|
||||
super(JointVelocityLowPassAction, self).process_actions(filtered)
|
||||
|
||||
def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(self.num_envs, device=self.device)
|
||||
elif isinstance(env_ids, slice):
|
||||
env_ids = torch.arange(self.num_envs, device=self.device)[env_ids]
|
||||
|
||||
self.reset_delays(env_ids)
|
||||
self._action_buffer[:, env_ids] = 0.0
|
||||
self._prev_raw[env_ids] = 0.0
|
||||
super().reset(env_ids)
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"""Discrete mode command term for multi-conditioned policy training.
|
||||
|
||||
Architecture is modeled after UniformVelocityCommand / UniformVelocityCommandCfg in
|
||||
mjlab/tasks/velocity/mdp/velocity_command.py.
|
||||
|
||||
During training: mode_id in {0, 1, 2} is sampled randomly per episode based on mode_probs.
|
||||
During deployment: ModeCommandCfg.fixed_mode is set to 0, 1, or 2, which can be modified externally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
|
||||
|
||||
|
||||
# Locomotion Mode ID Constants
|
||||
MODE_WALK = 0
|
||||
MODE_CLIMB = 1
|
||||
MODE_CROUCH = 2
|
||||
NUM_MODES = 3
|
||||
|
||||
|
||||
class ModeCommand(CommandTerm):
|
||||
"""Discrete locomotion mode command term.
|
||||
|
||||
The command tensor shape is [num_envs, 1], holding normalized mode values in [-1.0, 1.0].
|
||||
Mapping details: mode_id -> observation value
|
||||
0 (walk) -> -1.0
|
||||
1 (climb) -> 0.0
|
||||
2 (crouch) -> +1.0
|
||||
Using normalized scalar values keeps observation dimensions minimal (1D).
|
||||
"""
|
||||
|
||||
cfg: ModeCommandCfg
|
||||
|
||||
def __init__(self, cfg: ModeCommandCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
# Stores current integer mode_id for each environment
|
||||
self._mode_id = torch.zeros(
|
||||
self.num_envs, dtype=torch.long, device=self.device
|
||||
)
|
||||
# Normalized command observations: shape [num_envs, 1]
|
||||
self._command = torch.zeros(self.num_envs, 1, device=self.device)
|
||||
|
||||
# Pre-calculate cumulative sum of mode probabilities (normalized for multinomial sampling)
|
||||
probs = torch.tensor(cfg.mode_probs, dtype=torch.float32, device=self.device)
|
||||
self._mode_probs = probs / probs.sum()
|
||||
|
||||
@property
|
||||
def command(self) -> torch.Tensor:
|
||||
"""Normalized mode command observation tensor of shape [num_envs, 1]."""
|
||||
return self._command
|
||||
|
||||
def _update_metrics(self) -> None:
|
||||
pass
|
||||
|
||||
def _resample_command(self, env_ids: torch.Tensor) -> None:
|
||||
if self.cfg.fixed_mode is not None:
|
||||
# Deployment mode: fixed mode
|
||||
self._mode_id[env_ids] = self.cfg.fixed_mode
|
||||
else:
|
||||
# Training mode: map modes based on active terrain type if configured
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
terrain_types = getattr(terrain, "terrain_types", None)
|
||||
if self.cfg.terrain_mode_mapping is not None and terrain_types is not None:
|
||||
mapping = torch.tensor(
|
||||
self.cfg.terrain_mode_mapping, dtype=torch.long, device=self.device
|
||||
)
|
||||
types = terrain_types[env_ids]
|
||||
# Defensive clamp to prevent out of bounds indexing
|
||||
types_safe = torch.clamp(types, 0, len(mapping) - 1)
|
||||
self._mode_id[env_ids] = mapping[types_safe]
|
||||
else:
|
||||
# Fallback to random multinomial sampling
|
||||
sampled = torch.multinomial(
|
||||
self._mode_probs.expand(len(env_ids), -1),
|
||||
num_samples=1,
|
||||
replacement=True,
|
||||
).squeeze(-1)
|
||||
self._mode_id[env_ids] = sampled
|
||||
|
||||
# Update normalized command observations
|
||||
self._update_obs_from_mode_id(env_ids)
|
||||
|
||||
def _update_command(self) -> None:
|
||||
# Modes are persistent across episodes, no step-level command updates required
|
||||
pass
|
||||
|
||||
def _update_obs_from_mode_id(self, env_ids: torch.Tensor) -> None:
|
||||
"""Map integer mode_id to normalized observation: walk=0 -> -1.0, climb=1 -> 0.0, crouch=2 -> +1.0."""
|
||||
mode = self._mode_id[env_ids].float()
|
||||
# Linear projection: [0, NUM_MODES - 1] -> [-1.0, +1.0]
|
||||
normalized = 2.0 * mode / (NUM_MODES - 1) - 1.0
|
||||
self._command[env_ids, 0] = normalized
|
||||
|
||||
def create_gui(
|
||||
self,
|
||||
name: str,
|
||||
server, # ViserServer instance (late-bound to avoid top-level dependency)
|
||||
get_env_idx,
|
||||
on_change=None,
|
||||
request_action=None,
|
||||
) -> None:
|
||||
"""Construct interactive mode selection controls in the Viser visualizer GUI.
|
||||
|
||||
When enabled is active, slider values override target environment mode IDs in real-time.
|
||||
Discrete modes: 0=walk, 1=climb, 2=crouch (step size = 1).
|
||||
"""
|
||||
with server.gui.add_folder(f"{name.capitalize()} (0=walk 1=climb 2=crouch)"):
|
||||
enabled = server.gui.add_checkbox("Enable", initial_value=False)
|
||||
mode_slider = server.gui.add_slider(
|
||||
"mode_id",
|
||||
min=0,
|
||||
max=NUM_MODES - 1,
|
||||
step=1,
|
||||
initial_value=0,
|
||||
)
|
||||
|
||||
# Store GUI handles for compute() execution
|
||||
self._gui_enabled = enabled
|
||||
self._gui_slider = mode_slider
|
||||
self._gui_get_env_idx = get_env_idx
|
||||
|
||||
def compute(self, dt: float) -> None:
|
||||
"""Perform step updates: override mode_id with interactive sliders if visualizer GUI is active."""
|
||||
super().compute(dt)
|
||||
if (
|
||||
getattr(self, "_gui_enabled", None) is not None
|
||||
and self._gui_enabled.value
|
||||
and self._gui_get_env_idx is not None
|
||||
):
|
||||
idx = self._gui_get_env_idx()
|
||||
new_mode = int(round(self._gui_slider.value))
|
||||
env_ids = torch.tensor([idx], dtype=torch.long, device=self.device)
|
||||
self._mode_id[env_ids] = new_mode
|
||||
self._update_obs_from_mode_id(env_ids)
|
||||
|
||||
def get_mode_id(self, env_ids: torch.Tensor | None = None) -> torch.Tensor:
|
||||
"""Retrieve active environment mode IDs for reward/termination conditioning."""
|
||||
if env_ids is None:
|
||||
return self._mode_id
|
||||
return self._mode_id[env_ids]
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ModeCommandCfg(CommandTermCfg):
|
||||
"""Discrete ModeCommand configuration class."""
|
||||
|
||||
resampling_time_range: tuple[float, float] = (20.0, 20.0)
|
||||
"""Discrete mode resampling duration (resampled at episode boundary, equal to episode length)."""
|
||||
|
||||
mode_probs: tuple[float, ...] = (0.5, 0.3, 0.2)
|
||||
"""Probabilities for discrete modes (walk, climb, crouch). Normalized automatically."""
|
||||
|
||||
fixed_mode: int | None = None
|
||||
"""None for active random sampling during training; 0/1/2 for deployment."""
|
||||
|
||||
terrain_mode_mapping: tuple[int, ...] | None = None
|
||||
"""Force maps discrete modes based on environment sub-terrain grid index.
|
||||
|
||||
E.g., (0, 0, 0, 0, 1, 1, 2) corresponds to walk mode for first four sub-terrains,
|
||||
climb mode for the next two sub-terrains, and crouch mode for the last sub-terrain.
|
||||
Overrides mode_probs if active.
|
||||
"""
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> ModeCommand:
|
||||
return ModeCommand(self, env)
|
||||
@@ -0,0 +1,42 @@
|
||||
"""HIMLoco-style only_positive_rewards: clip total step reward to >= 0.
|
||||
|
||||
Mechanics:
|
||||
- PPO Advantage = reward + gamma * V(s') - V(s)
|
||||
- During early policy learning, when penalties greatly exceed positive rewards,
|
||||
the total step reward can become highly negative, producing negative advantages.
|
||||
- This often causes policies to learn to intentionally "fall down immediately" to terminate
|
||||
the episode and minimize long-term negative accumulation.
|
||||
- Clamping the total step reward to >= 0 prevents the penalty from generating negative
|
||||
step feedback, ensuring penalty terms can only offset positive gains without directing the
|
||||
advantage signal incorrectly.
|
||||
- Individual term episode sums are still recorded correctly for tracking and diagnostics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from mjlab.managers.reward_manager import RewardManager
|
||||
|
||||
_original_compute = RewardManager.compute
|
||||
|
||||
|
||||
def _compute_with_clamp(self: RewardManager, dt: float) -> torch.Tensor:
|
||||
"""Compute step rewards and clamp the final total reward to >= 0."""
|
||||
reward = _original_compute(self, dt)
|
||||
# In-place clamp on the same buffer reference for consistency with original API
|
||||
return torch.clamp(reward, min=0.0)
|
||||
|
||||
|
||||
def enable_only_positive_rewards() -> None:
|
||||
"""Enable HIMLoco-style positive-only step reward clamping globally.
|
||||
|
||||
Idempotent: calling this multiple times is completely safe.
|
||||
"""
|
||||
if RewardManager.compute is not _compute_with_clamp:
|
||||
RewardManager.compute = _compute_with_clamp
|
||||
|
||||
|
||||
def disable_only_positive_rewards() -> None:
|
||||
"""Disable positive-only step reward clamping, restoring default behavior."""
|
||||
if RewardManager.compute is not _original_compute:
|
||||
RewardManager.compute = _original_compute
|
||||
@@ -0,0 +1,738 @@
|
||||
"""MDP Reward functions for unitree locomotion training."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
import torch
|
||||
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.utils.lab_api.string import resolve_matching_names_values
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
def track_linear_velocity(
|
||||
env: ManagerBasedRlEnv,
|
||||
std: float,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reward tracking commanded horizontal xy velocity."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
actual = asset.data.root_link_lin_vel_b
|
||||
xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
|
||||
reward = torch.exp(-xy_error / std**2)
|
||||
return reward
|
||||
|
||||
|
||||
def track_angular_velocity(
|
||||
env: ManagerBasedRlEnv,
|
||||
std: float,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reward tracking commanded yaw rate."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
actual = asset.data.root_link_ang_vel_b
|
||||
z_error = torch.square(command[:, 2] - actual[:, 2])
|
||||
reward = torch.exp(-z_error / std**2)
|
||||
return reward
|
||||
|
||||
|
||||
def base_height_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
target_height: float = 0.36,
|
||||
sensor_cfg: SceneEntityCfg | None = None,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize base height deviation from target height, relative to terrain or origin."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
has_sensor = False
|
||||
if sensor_cfg is not None:
|
||||
try:
|
||||
_ = env.scene[sensor_cfg.name]
|
||||
has_sensor = True
|
||||
except KeyError:
|
||||
has_sensor = False
|
||||
|
||||
if has_sensor:
|
||||
from mjlab.envs.mdp.observations import height_scan
|
||||
# Get height scanner readings (each ray's height relative to base frame Z)
|
||||
heights = height_scan(env, sensor_cfg.name, offset=0.0)
|
||||
# Clean any invalid values
|
||||
heights = torch.nan_to_num(heights, nan=target_height, posinf=target_height, neginf=target_height)
|
||||
|
||||
# Filter out rays that represent extreme misses/cliffs (e.g. height < 0 or height > 1.5)
|
||||
# The robot target height is ~0.36m, so the scanner should measure something between 0.1m and 0.8m usually.
|
||||
valid_mask = (heights > 0.0) & (heights < 1.5)
|
||||
|
||||
# Calculate mean height above terrain ignoring misses
|
||||
sum_heights = torch.sum(torch.where(valid_mask, heights, torch.zeros_like(heights)), dim=1)
|
||||
count_valid = torch.sum(valid_mask.float(), dim=1)
|
||||
measured_height = torch.where(count_valid > 0, sum_heights / torch.clamp(count_valid, min=1.0), torch.full_like(sum_heights, target_height))
|
||||
|
||||
error = measured_height - target_height
|
||||
else:
|
||||
if hasattr(env.scene, "env_origins") and env.scene.env_origins is not None:
|
||||
env_origins_z = env.scene.env_origins[:, 2]
|
||||
else:
|
||||
env_origins_z = torch.zeros(env.num_envs, device=env.device)
|
||||
|
||||
root_z = asset.data.root_link_pos_w[:, 2] - env_origins_z
|
||||
error = root_z - target_height
|
||||
|
||||
reward = torch.square(error)
|
||||
return reward
|
||||
|
||||
|
||||
def safe_height_scan(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
"""Safely fetch height scanner outputs, replacing NaNs/infs with finite values."""
|
||||
from mjlab.envs.mdp.observations import height_scan
|
||||
result = height_scan(env, sensor_name)
|
||||
return torch.nan_to_num(result, nan=0.0, posinf=5.0, neginf=-5.0)
|
||||
|
||||
|
||||
def safe_base_lin_vel(env: ManagerBasedRlEnv) -> torch.Tensor:
|
||||
"""Safely fetch base linear velocity, replacing NaNs/infs with finite values."""
|
||||
from mjlab.envs.mdp.observations import base_lin_vel
|
||||
result = base_lin_vel(env)
|
||||
return torch.nan_to_num(result, nan=0.0, posinf=100.0, neginf=-100.0)
|
||||
|
||||
|
||||
def safe_foot_air_time(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
"""Safely fetch foot air time, replacing NaNs with zeros."""
|
||||
from mjlab.tasks.velocity.mdp.observations import foot_air_time
|
||||
result = foot_air_time(env, sensor_name)
|
||||
return torch.nan_to_num(result, nan=0.0)
|
||||
|
||||
|
||||
def safe_foot_contact(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
"""Safely fetch foot contact flags, replacing NaNs with zeros."""
|
||||
from mjlab.tasks.velocity.mdp.observations import foot_contact
|
||||
result = foot_contact(env, sensor_name)
|
||||
return torch.nan_to_num(result, nan=0.0)
|
||||
|
||||
|
||||
def safe_foot_contact_forces(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
"""Safely fetch foot contact forces, replacing NaNs with zeros."""
|
||||
from mjlab.tasks.velocity.mdp.observations import foot_contact_forces
|
||||
result = foot_contact_forces(env, sensor_name)
|
||||
return torch.nan_to_num(result, nan=0.0)
|
||||
|
||||
|
||||
def feet_clearance(
|
||||
env: ManagerBasedRlEnv,
|
||||
target_height: float,
|
||||
command_name: str = "twist",
|
||||
command_threshold: float = 0.1,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize deviation from the target foot clearance height."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", body_names=(".*_wheel_Link",))
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
foot_z = asset.data.body_com_pos_w[:, asset_cfg.body_ids, 2]
|
||||
foot_vel_xy = asset.data.body_link_vel_w[:, asset_cfg.body_ids, :2]
|
||||
vel_norm = torch.norm(foot_vel_xy, dim=-1)
|
||||
delta = torch.abs(foot_z - target_height)
|
||||
cost = torch.sum(delta * vel_norm, dim=1)
|
||||
|
||||
cmd = env.command_manager.get_command(command_name)
|
||||
linear_norm = torch.norm(cmd[:, :2], dim=1)
|
||||
angular_norm = torch.abs(cmd[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
|
||||
reward = cost * active
|
||||
return reward
|
||||
|
||||
|
||||
def soft_landing(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
command_name: str = "twist",
|
||||
command_threshold: float = 0.05,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize high impact forces on the first landing contact."""
|
||||
from mjlab.sensor import ContactSensor
|
||||
contact_sensor: ContactSensor = env.scene[sensor_name]
|
||||
forces = contact_sensor.data.force
|
||||
force_magnitude = torch.norm(forces, dim=-1)
|
||||
first_contact = contact_sensor.compute_first_contact(dt=env.step_dt)
|
||||
landing_impact = force_magnitude * first_contact.float()
|
||||
cost = torch.sum(landing_impact, dim=1)
|
||||
|
||||
cmd = env.command_manager.get_command(command_name)
|
||||
linear_norm = torch.norm(cmd[:, :2], dim=1)
|
||||
angular_norm = torch.abs(cmd[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
|
||||
reward = cost * active
|
||||
return reward
|
||||
|
||||
|
||||
def _total_command_magnitude(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
|
||||
"""Retrieve absolute horizontal command speed."""
|
||||
command = env.command_manager.get_command(command_name)
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
return linear_norm + angular_norm
|
||||
|
||||
|
||||
def wheel_roll_tracking(
|
||||
env: ManagerBasedRlEnv,
|
||||
command_name: str,
|
||||
wheel_radius: float,
|
||||
wheel_track: float,
|
||||
std: float,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reward wheel angular velocity matching the commanded planar motion."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
|
||||
wheel_vel = asset.data.joint_vel[:, asset_cfg.joint_ids]
|
||||
# Assumes order: fl, fr, rl, rr. Side signs: fl=1, fr=-1, rl=1, rr=-1
|
||||
side_signs = torch.tensor(
|
||||
[1.0, -1.0, 1.0, -1.0],
|
||||
device=env.device,
|
||||
dtype=wheel_vel.dtype,
|
||||
).unsqueeze(0)
|
||||
|
||||
lin_vel_x = command[:, 0:1]
|
||||
ang_vel_z = command[:, 2:3]
|
||||
target_wheel_vel = (
|
||||
lin_vel_x + 0.5 * wheel_track * side_signs * ang_vel_z
|
||||
) / wheel_radius
|
||||
|
||||
err = torch.mean(torch.square(wheel_vel - target_wheel_vel), dim=1)
|
||||
reward = torch.exp(-err / (std**2))
|
||||
return reward
|
||||
|
||||
|
||||
def adaptive_leg_motion_penalty(
|
||||
env: ManagerBasedRlEnv,
|
||||
command_name: str,
|
||||
sensor_name: str,
|
||||
command_threshold: float = 0.05,
|
||||
tilt_relax_start: float = 0.08,
|
||||
tilt_relax_end: float = 0.30,
|
||||
contact_target: float = 0.85,
|
||||
min_penalty_scale: float = 0.2,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize leg motion less when tilt or wheel contact suggests recovery."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", joint_names=(
|
||||
".*_hip_abduction_joint",
|
||||
".*_hip_pitch_joint",
|
||||
".*_knee_joint",
|
||||
))
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
leg_vel = asset.data.joint_vel[:, asset_cfg.joint_ids]
|
||||
base_cost = torch.mean(torch.square(leg_vel), dim=1)
|
||||
active = (_total_command_magnitude(env, command_name) > command_threshold).float()
|
||||
|
||||
tilt_xy = torch.norm(asset.data.projected_gravity_b[:, :2], dim=1)
|
||||
tilt_relax = torch.clamp(
|
||||
(tilt_xy - tilt_relax_start) / max(tilt_relax_end - tilt_relax_start, 1.0e-6),
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
)
|
||||
|
||||
from mjlab.sensor import ContactSensor
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
in_contact = (sensor.data.found > 0).float()
|
||||
contact_fraction = torch.mean(in_contact, dim=1)
|
||||
if contact_fraction.ndim > 1:
|
||||
contact_fraction = torch.mean(contact_fraction, dim=1)
|
||||
contact_relax = torch.clamp(
|
||||
(contact_target - contact_fraction) / max(contact_target, 1.0e-6),
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
)
|
||||
|
||||
relax = torch.maximum(tilt_relax, contact_relax)
|
||||
penalty_scale = 1.0 - (1.0 - min_penalty_scale) * relax
|
||||
reward = base_cost * active * penalty_scale
|
||||
return reward
|
||||
|
||||
|
||||
def contact_fraction_reward(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
) -> torch.Tensor:
|
||||
"""Reward persistent wheel-ground contact."""
|
||||
from mjlab.sensor import ContactSensor
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
in_contact = (sensor.data.found > 0).float()
|
||||
reward = torch.mean(in_contact, dim=1)
|
||||
return reward
|
||||
|
||||
|
||||
def stand_still(
|
||||
env,
|
||||
command_name: str,
|
||||
command_threshold: float = 0.1,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", joint_names=(
|
||||
".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint",
|
||||
))
|
||||
asset = env.scene[asset_cfg.name]
|
||||
diff = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
cost = torch.mean(torch.abs(diff), dim=1)
|
||||
command = env.command_manager.get_command(command_name)
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
inactive = (linear_norm + angular_norm < command_threshold).float()
|
||||
reward = cost * inactive
|
||||
return reward
|
||||
|
||||
def hip_deviation(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize hip abduction joints deviating from default."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint",))
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
q = asset.data.joint_pos[:, asset_cfg.joint_ids]
|
||||
q0 = asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
reward = torch.sum(torch.square(q - q0), dim=1)
|
||||
return reward
|
||||
|
||||
|
||||
def joint_deviation_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize all specified joints deviating from default."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
q = asset.data.joint_pos[:, asset_cfg.joint_ids]
|
||||
q0 = asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
reward = torch.sum(torch.square(q - q0), dim=1)
|
||||
return reward
|
||||
|
||||
|
||||
def flat_orientation_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize non-flat base orientation (roll/pitch via projected gravity)."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
reward = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
|
||||
return reward
|
||||
|
||||
|
||||
def lin_vel_z_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize vertical (z) base linear velocity."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
reward = torch.square(asset.data.root_link_lin_vel_b[:, 2])
|
||||
return reward
|
||||
|
||||
|
||||
class variable_posture:
|
||||
"""Reward posture tracking with speed-dependent tolerance."""
|
||||
|
||||
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
|
||||
asset: Entity = env.scene[cfg.params["asset_cfg"].name]
|
||||
default_joint_pos = asset.data.default_joint_pos
|
||||
self.default_joint_pos = default_joint_pos
|
||||
|
||||
_, joint_names = asset.find_joints(cfg.params["asset_cfg"].joint_names)
|
||||
|
||||
_, _, std_standing = resolve_matching_names_values(
|
||||
data=cfg.params["std_standing"], list_of_strings=joint_names,
|
||||
)
|
||||
self.std_standing = torch.tensor(std_standing, device=env.device, dtype=torch.float32)
|
||||
|
||||
_, _, std_walking = resolve_matching_names_values(
|
||||
data=cfg.params["std_walking"], list_of_strings=joint_names,
|
||||
)
|
||||
self.std_walking = torch.tensor(std_walking, device=env.device, dtype=torch.float32)
|
||||
|
||||
_, _, std_running = resolve_matching_names_values(
|
||||
data=cfg.params["std_running"], list_of_strings=joint_names,
|
||||
)
|
||||
self.std_running = torch.tensor(std_running, device=env.device, dtype=torch.float32)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
std_standing: dict,
|
||||
std_walking: dict,
|
||||
std_running: dict,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
command_name: str,
|
||||
walking_threshold: float = 0.5,
|
||||
running_threshold: float = 1.5,
|
||||
) -> torch.Tensor:
|
||||
del std_standing, std_walking, std_running
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
total_speed = _total_command_magnitude(env, command_name)
|
||||
|
||||
# 提取身体倾斜度 (重力在 XY 平面的投影)
|
||||
tilt_xy = torch.norm(asset.data.projected_gravity_b[:, :2], dim=1)
|
||||
|
||||
# 如果身体倾斜 > 0.15 (约 8.6 度)，说明正在爬墙、越障或已摔倒
|
||||
is_tilting = (tilt_xy > 0.15).float()
|
||||
is_running_cmd = (total_speed >= running_threshold).float()
|
||||
|
||||
# 【核心越障修改】只要速度极快，或者正在倾斜攀爬，直接赋予最大幅度的动作容忍权 (running_mask)
|
||||
running_mask = torch.clamp(is_running_cmd + is_tilting, 0.0, 1.0)
|
||||
|
||||
# 剩下的平稳且不倾斜的状态，再去根据低速指令判断是 walking 还是 standing
|
||||
not_running_mask = 1.0 - running_mask
|
||||
is_walking_cmd = (total_speed >= walking_threshold).float()
|
||||
|
||||
walking_mask = not_running_mask * is_walking_cmd
|
||||
standing_mask = not_running_mask * (1.0 - is_walking_cmd)
|
||||
|
||||
std = (
|
||||
self.std_standing * standing_mask.unsqueeze(1)
|
||||
+ self.std_walking * walking_mask.unsqueeze(1)
|
||||
+ self.std_running * running_mask.unsqueeze(1)
|
||||
)
|
||||
|
||||
current_joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids]
|
||||
desired_joint_pos = self.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
error_squared = torch.square(current_joint_pos - desired_joint_pos)
|
||||
reward = torch.exp(-torch.mean(error_squared / (std**2), dim=1))
|
||||
return reward
|
||||
|
||||
|
||||
def get_mode_id(env: ManagerBasedRlEnv) -> torch.Tensor:
|
||||
"""Retrieve current integer mode_id tensor from command_manager.
|
||||
|
||||
Safely falls back to zeros (walk mode) if command manager or mode is not initialized.
|
||||
"""
|
||||
if not hasattr(env, "command_manager") or env.command_manager is None:
|
||||
return torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
|
||||
term = env.command_manager.get_term("mode")
|
||||
if term is None or not hasattr(term, "get_mode_id"):
|
||||
return torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
|
||||
return term.get_mode_id()
|
||||
|
||||
|
||||
def crouch_height_reward(
|
||||
env: ManagerBasedRlEnv,
|
||||
target_height: float = 0.22,
|
||||
std: float = 0.05,
|
||||
mode_command_name: str = "mode",
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reward the robot for keeping body low under target height when in crouch mode (mode_id == 2).
|
||||
|
||||
Uses a single-sided penalty so that going lower than target_height receives full reward.
|
||||
"""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
if hasattr(env.scene, "env_origins") and env.scene.env_origins is not None:
|
||||
env_origins_z = env.scene.env_origins[:, 2]
|
||||
else:
|
||||
env_origins_z = torch.zeros(env.num_envs, device=env.device)
|
||||
|
||||
root_z = asset.data.root_link_pos_w[:, 2] - env_origins_z
|
||||
error = torch.square(torch.clamp(root_z - target_height, min=0.0))
|
||||
reward = torch.exp(-error / std ** 2)
|
||||
|
||||
mode_id = get_mode_id(env)
|
||||
active = (mode_id == 2).float()
|
||||
reward = reward * active
|
||||
return reward
|
||||
|
||||
|
||||
class mode_conditioned_posture:
|
||||
"""Conditioned posture reward selecting targets/tolerances by active mode_id."""
|
||||
|
||||
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv) -> None:
|
||||
asset = env.scene[cfg.params["asset_cfg"].name]
|
||||
default_joint_pos = asset.data.default_joint_pos
|
||||
assert default_joint_pos is not None
|
||||
self.default_joint_pos = default_joint_pos
|
||||
|
||||
_, joint_names = asset.find_joints(cfg.params["asset_cfg"].joint_names)
|
||||
self.joint_ids = asset.find_joints(cfg.params["asset_cfg"].joint_names)[0]
|
||||
|
||||
_, _, std_walk = resolve_matching_names_values(
|
||||
data=cfg.params["std_walk"], list_of_strings=joint_names,
|
||||
)
|
||||
self.std_walk = torch.tensor(std_walk, device=env.device, dtype=torch.float32)
|
||||
|
||||
if "crouch_pos" in cfg.params:
|
||||
_, _, crouch_pos_vals = resolve_matching_names_values(
|
||||
data=cfg.params["crouch_pos"], list_of_strings=joint_names,
|
||||
)
|
||||
self.crouch_pos_target = torch.tensor(
|
||||
crouch_pos_vals, device=env.device, dtype=torch.float32
|
||||
).unsqueeze(0)
|
||||
else:
|
||||
self.crouch_pos_target = None
|
||||
|
||||
climb_scale = cfg.params.get("climb_std_scale", 3.0)
|
||||
crouch_scale = cfg.params.get("crouch_std_scale", 4.0)
|
||||
self.std_climb = self.std_walk * climb_scale
|
||||
self.std_crouch = self.std_walk * crouch_scale
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
std_walk: dict,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
climb_std_scale: float = 3.0,
|
||||
crouch_std_scale: float = 4.0,
|
||||
crouch_pos: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
del std_walk, crouch_pos
|
||||
|
||||
asset = env.scene[asset_cfg.name]
|
||||
mode_id = get_mode_id(env).float()
|
||||
|
||||
walk_mask = (mode_id == 0).float().unsqueeze(1)
|
||||
climb_mask = (mode_id == 1).float().unsqueeze(1)
|
||||
crouch_mask = (mode_id == 2).float().unsqueeze(1)
|
||||
|
||||
std = (
|
||||
self.std_walk * walk_mask
|
||||
+ self.std_climb * climb_mask
|
||||
+ self.std_crouch * crouch_mask
|
||||
)
|
||||
|
||||
current_pos = asset.data.joint_pos[:, self.joint_ids]
|
||||
base_desired_pos = self.default_joint_pos[:, self.joint_ids]
|
||||
|
||||
if self.crouch_pos_target is not None:
|
||||
desired_pos = (
|
||||
base_desired_pos * (walk_mask + climb_mask)
|
||||
+ self.crouch_pos_target * crouch_mask
|
||||
)
|
||||
else:
|
||||
desired_pos = base_desired_pos
|
||||
|
||||
error_sq = torch.square(current_pos - desired_pos)
|
||||
reward = torch.exp(-torch.mean(error_sq / (std ** 2 + 1e-8), dim=1))
|
||||
return reward
|
||||
|
||||
|
||||
def crawl_height_reward(
|
||||
env: ManagerBasedRlEnv,
|
||||
target_height: float = 0.22,
|
||||
std: float = 0.05,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reward the robot for staying below target crawling height (0.22m).
|
||||
|
||||
Uses a single-sided penalty: error is 0 if root_z <= target_height.
|
||||
"""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
if hasattr(env.scene, "env_origins") and env.scene.env_origins is not None:
|
||||
env_origins_z = env.scene.env_origins[:, 2]
|
||||
else:
|
||||
env_origins_z = torch.zeros(env.num_envs, device=env.device)
|
||||
|
||||
root_z = asset.data.root_link_pos_w[:, 2] - env_origins_z
|
||||
error = torch.square(torch.clamp(root_z - target_height, min=0.0))
|
||||
reward = torch.exp(-error / std ** 2)
|
||||
return reward
|
||||
|
||||
|
||||
def _get_terrain_levels(env: ManagerBasedRlEnv) -> torch.Tensor | None:
|
||||
"""Safely fetch current active terrain levels from environment."""
|
||||
terrain = getattr(env.scene, "terrain", None)
|
||||
if terrain is not None and hasattr(terrain, "terrain_levels"):
|
||||
return terrain.terrain_levels
|
||||
if hasattr(env, "terrain_levels"):
|
||||
return env.terrain_levels
|
||||
return None
|
||||
|
||||
|
||||
def terrain_level_bonus(env: ManagerBasedRlEnv) -> torch.Tensor:
|
||||
"""Positive step bonus normalized to [0, 1] by max terrain level (num_rows-1 = 9).
|
||||
|
||||
Normalizing prevents the raw level integer (0~9) from dominating the reward
|
||||
signal at high difficulty, which would otherwise incentivize level-rushing
|
||||
over quality locomotion.
|
||||
"""
|
||||
levels = _get_terrain_levels(env)
|
||||
if levels is None:
|
||||
return torch.zeros(env.num_envs, device=env.device)
|
||||
return levels.float() / 9.0
|
||||
|
||||
|
||||
def action_rate_curriculum_l2(env: ManagerBasedRlEnv) -> torch.Tensor:
|
||||
"""Apply action change rate penalty with curriculum-driven scaling.
|
||||
|
||||
Penalizes rate changes fully at level 0 (flat) for optimal smoothness,
|
||||
decaying to 10% penalty at levels 8+ to facilitate explosive dynamic clearing maneuvers.
|
||||
"""
|
||||
action_rate = torch.sum(torch.square(env.action_manager.action - env.action_manager.prev_action), dim=1)
|
||||
levels = _get_terrain_levels(env)
|
||||
if levels is None:
|
||||
reward = action_rate
|
||||
else:
|
||||
decay = 1.0 - 0.9 * torch.clamp(levels.float() / 8.0, 0.0, 1.0)
|
||||
reward = action_rate * decay
|
||||
|
||||
return reward
|
||||
|
||||
|
||||
def leg_symmetry(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize asymmetry between left and right legs (pitch and knee joints) robustly."""
|
||||
asset: Entity = env.scene["robot"]
|
||||
joint_names = asset.data.joint_names
|
||||
|
||||
def find_idx(side, jtype):
|
||||
for i, name in enumerate(joint_names):
|
||||
name_lower = name.lower()
|
||||
if side in name_lower and jtype in name_lower:
|
||||
return i
|
||||
# Fallback if not found to avoid crash, though it should find it
|
||||
return 0
|
||||
|
||||
fl_p = find_idx("fl_", "pitch")
|
||||
fr_p = find_idx("fr_", "pitch")
|
||||
rl_p = find_idx("rl_", "pitch")
|
||||
rr_p = find_idx("rr_", "pitch")
|
||||
|
||||
fl_k = find_idx("fl_", "knee")
|
||||
fr_k = find_idx("fr_", "knee")
|
||||
rl_k = find_idx("rl_", "knee")
|
||||
rr_k = find_idx("rr_", "knee")
|
||||
|
||||
q = asset.data.joint_pos
|
||||
cost = torch.square(q[:, fl_p] - q[:, fr_p]) + \
|
||||
torch.square(q[:, rl_p] - q[:, rr_p]) + \
|
||||
torch.square(q[:, fl_k] - q[:, fr_k]) + \
|
||||
torch.square(q[:, rl_k] - q[:, rr_k])
|
||||
return cost
|
||||
|
||||
def feet_contact_without_cmd(env, command_name: str, sensor_name: str) -> torch.Tensor:
|
||||
from mjlab.sensor import ContactSensor
|
||||
contact_sensor = env.scene[sensor_name]
|
||||
contact = contact_sensor.data.found > 0
|
||||
reward = torch.sum(contact, dim=-1).float()
|
||||
cmd = env.command_manager.get_command(command_name)
|
||||
linear_norm = torch.norm(cmd[:, :2], dim=1)
|
||||
angular_norm = torch.abs(cmd[:, 2])
|
||||
reward *= (linear_norm + angular_norm) < 0.1
|
||||
return reward
|
||||
|
||||
def joint_pos_penalty(
|
||||
env,
|
||||
command_name: str,
|
||||
stand_still_scale: float,
|
||||
velocity_threshold: float,
|
||||
command_threshold: float,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset = env.scene[asset_cfg.name]
|
||||
cmd_val = env.command_manager.get_command(command_name)
|
||||
cmd = torch.linalg.norm(cmd_val[:, :2], dim=1) + torch.abs(cmd_val[:, 2])
|
||||
body_vel = torch.linalg.norm(asset.data.root_link_lin_vel_b[:, :2], dim=1)
|
||||
running_reward = torch.linalg.norm(
|
||||
(asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
|
||||
)
|
||||
reward = torch.where(
|
||||
torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold),
|
||||
running_reward,
|
||||
stand_still_scale * running_reward,
|
||||
)
|
||||
return reward
|
||||
|
||||
def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset = env.scene[asset_cfg.name]
|
||||
if not hasattr(env, "joint_mirror_joints_cache") or env.joint_mirror_joints_cache is None:
|
||||
env.joint_mirror_joints_cache = []
|
||||
for joint_pair in mirror_joints:
|
||||
p0 = asset.find_joints(joint_pair[0])[0]
|
||||
p1 = asset.find_joints(joint_pair[1])[0]
|
||||
env.joint_mirror_joints_cache.append([p0, p1])
|
||||
reward = torch.zeros(env.num_envs, device=env.device)
|
||||
for joint_pair in env.joint_mirror_joints_cache:
|
||||
diff = torch.sum(
|
||||
torch.square(asset.data.joint_pos[:, joint_pair[0]] - asset.data.joint_pos[:, joint_pair[1]]),
|
||||
dim=-1,
|
||||
)
|
||||
reward += diff
|
||||
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
|
||||
return reward
|
||||
|
||||
|
||||
|
||||
def upward(env, asset_cfg=None):
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg('robot')
|
||||
asset = env.scene[asset_cfg.name]
|
||||
reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
|
||||
return reward
|
||||
|
||||
def upright_roll_only(env, asset_cfg=None):
|
||||
if asset_cfg is None:
|
||||
from mjlab.envs.manager_based_rl_env import SceneEntityCfg
|
||||
asset_cfg = SceneEntityCfg('robot')
|
||||
asset = env.scene[asset_cfg.name]
|
||||
reward = torch.square(asset.data.projected_gravity_b[:, 1])
|
||||
return reward
|
||||
|
||||
def track_linear_velocity_l1(env, std, command_name):
|
||||
asset = env.scene['robot']
|
||||
cmd = env.command_manager.get_command(command_name)
|
||||
lin_vel_error = torch.linalg.norm(asset.data.root_link_lin_vel_b[:, :2] - cmd[:, :2], dim=1)
|
||||
reward = torch.exp(-lin_vel_error / std)
|
||||
return reward
|
||||
|
||||
def pitch_control_penalty(env, max_pitch_rad: float = 0.50, asset_cfg=None) -> torch.Tensor:
|
||||
"""Penalize excessive pitch orientation beyond a safe threshold.
|
||||
|
||||
Allows pitch angles up to max_pitch_rad (e.g. 29 degrees) for climbing obstacles,
|
||||
but quadratically penalizes any pitch angle exceeding this threshold (e.g. during flinging).
|
||||
"""
|
||||
import math
|
||||
if asset_cfg is None:
|
||||
from mjlab.envs.manager_based_rl_env import SceneEntityCfg
|
||||
asset_cfg = SceneEntityCfg('robot')
|
||||
asset = env.scene[asset_cfg.name]
|
||||
g_x = asset.data.projected_gravity_b[:, 0]
|
||||
g_x_threshold = math.sin(max_pitch_rad)
|
||||
excessive_pitch = torch.clamp(torch.abs(g_x) - g_x_threshold, min=0.0)
|
||||
reward = torch.square(excessive_pitch)
|
||||
return reward
|
||||
@@ -0,0 +1,41 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
"""
|
||||
HimLoco RSL-RL implementation with history-informed models.
|
||||
"""
|
||||
|
||||
# Export HIM implementations
|
||||
from .algorithms.him_ppo import *
|
||||
from .modules.him_actor_critic import *
|
||||
from .modules.him_estimator import *
|
||||
from .storage.him_rollout_storage import *
|
||||
from .runners.him_on_policy_runner import *
|
||||
from .wrappers import *
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .him_ppo import HIMPPO
|
||||
@@ -0,0 +1,192 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from ..modules import HIMActorCritic
|
||||
from ..storage import HIMRolloutStorage
|
||||
|
||||
class HIMPPO:
|
||||
actor_critic: HIMActorCritic
|
||||
def __init__(self,
|
||||
actor_critic,
|
||||
num_learning_epochs=1,
|
||||
num_mini_batches=1,
|
||||
clip_param=0.2,
|
||||
gamma=0.998,
|
||||
lam=0.95,
|
||||
value_loss_coef=1.0,
|
||||
entropy_coef=0.0,
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
device='cpu',
|
||||
):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
|
||||
# PPO components
|
||||
self.actor_critic = actor_critic
|
||||
self.actor_critic.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
|
||||
self.transition = HIMRolloutStorage.Transition()
|
||||
|
||||
# PPO parameters
|
||||
self.clip_param = clip_param
|
||||
self.num_learning_epochs = num_learning_epochs
|
||||
self.num_mini_batches = num_mini_batches
|
||||
self.value_loss_coef = value_loss_coef
|
||||
self.entropy_coef = entropy_coef
|
||||
self.gamma = gamma
|
||||
self.lam = lam
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.use_clipped_value_loss = use_clipped_value_loss
|
||||
|
||||
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
|
||||
self.storage = HIMRolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)
|
||||
|
||||
def test_mode(self):
|
||||
self.actor_critic.test()
|
||||
|
||||
def train_mode(self):
|
||||
self.actor_critic.train()
|
||||
|
||||
def act(self, obs, critic_obs):
|
||||
# Compute the actions and values
|
||||
self.transition.actions = self.actor_critic.act(obs).detach()
|
||||
self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
|
||||
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
|
||||
self.transition.action_mean = self.actor_critic.action_mean.detach()
|
||||
self.transition.action_sigma = self.actor_critic.action_std.detach()
|
||||
# need to record obs and critic_obs before env.step()
|
||||
self.transition.observations = obs
|
||||
self.transition.critic_observations = critic_obs
|
||||
return self.transition.actions
|
||||
|
||||
def process_env_step(self, rewards, dones, infos, next_critic_obs):
|
||||
self.transition.next_critic_observations = next_critic_obs.clone()
|
||||
self.transition.rewards = rewards.clone()
|
||||
self.transition.dones = dones
|
||||
# Bootstrapping on time outs
|
||||
if 'time_outs' in infos:
|
||||
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transitions(self.transition)
|
||||
self.transition.clear()
|
||||
self.actor_critic.reset(dones)
|
||||
|
||||
def compute_returns(self, last_critic_obs):
|
||||
last_values= self.actor_critic.evaluate(last_critic_obs).detach()
|
||||
self.storage.compute_returns(last_values, self.gamma, self.lam)
|
||||
|
||||
def update(self):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_estimation_loss = 0
|
||||
mean_swap_loss = 0
|
||||
|
||||
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
|
||||
for obs_batch, critic_obs_batch, actions_batch, next_critic_obs_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
|
||||
old_mu_batch, old_sigma_batch in generator:
|
||||
|
||||
self.actor_critic.act(obs_batch)
|
||||
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.actor_critic.evaluate(critic_obs_batch)
|
||||
mu_batch = self.actor_critic.action_mean
|
||||
sigma_batch = self.actor_critic.action_std
|
||||
entropy_batch = self.actor_critic.entropy
|
||||
|
||||
# KL
|
||||
if self.desired_kl != None and self.schedule == 'adaptive':
|
||||
with torch.inference_mode():
|
||||
kl = torch.sum(
|
||||
torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
|
||||
kl_mean = torch.mean(kl)
|
||||
|
||||
if kl_mean > self.desired_kl * 2.0:
|
||||
self.learning_rate = max(1e-5, self.learning_rate / 1.5)
|
||||
elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
|
||||
self.learning_rate = min(1e-2, self.learning_rate * 1.5)
|
||||
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
#Estimator Update
|
||||
estimation_loss, swap_loss = self.actor_critic.estimator.update(obs_batch, next_critic_obs_batch, lr=self.learning_rate)
|
||||
|
||||
# Surrogate loss
|
||||
ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
|
||||
surrogate = -torch.squeeze(advantages_batch) * ratio
|
||||
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
|
||||
1.0 + self.clip_param)
|
||||
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
|
||||
|
||||
# Value function loss
|
||||
if self.use_clipped_value_loss:
|
||||
value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
|
||||
self.clip_param)
|
||||
value_losses = (value_batch - returns_batch).pow(2)
|
||||
value_losses_clipped = (value_clipped - returns_batch).pow(2)
|
||||
value_loss = torch.max(value_losses, value_losses_clipped).mean()
|
||||
else:
|
||||
value_loss = (returns_batch - value_batch).pow(2).mean()
|
||||
|
||||
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
|
||||
|
||||
# Gradient step
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_estimation_loss += estimation_loss
|
||||
mean_swap_loss += swap_loss
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_estimation_loss /= num_updates
|
||||
mean_swap_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, estimation_loss, swap_loss
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .rl_cfg import *
|
||||
@@ -0,0 +1,180 @@
|
||||
from dataclasses import MISSING
|
||||
from isaaclab.utils import configclass
|
||||
from typing import Literal
|
||||
|
||||
|
||||
@configclass
|
||||
class HIMBaseRunnerCfg:
|
||||
"""Base configuration of the runner."""
|
||||
|
||||
seed: int = 1
|
||||
"""The seed for the experiment. Default is 1."""
|
||||
|
||||
device: str = "cuda:0"
|
||||
"""The device for the rl-agent. Default is cuda:0."""
|
||||
|
||||
num_steps_per_env: int = MISSING
|
||||
"""The number of steps per environment per update."""
|
||||
|
||||
max_iterations: int = MISSING
|
||||
"""The maximum number of iterations."""
|
||||
|
||||
empirical_normalization: bool | None = None
|
||||
"""This parameter is deprecated and will be removed in the future.
|
||||
|
||||
Use `actor_obs_normalization` and `critic_obs_normalization` instead.
|
||||
"""
|
||||
|
||||
# obs_groups: dict[str, list[str]] = MISSING
|
||||
# """A mapping from observation groups to observation sets.
|
||||
|
||||
# The keys of the dictionary are predefined observation sets used by the underlying algorithm
|
||||
# and values are lists of observation groups provided by the environment.
|
||||
|
||||
# For instance, if the environment provides a dictionary of observations with groups "policy", "images",
|
||||
# and "privileged", these can be mapped to algorithmic observation sets as follows:
|
||||
|
||||
# .. code-block:: python
|
||||
|
||||
# obs_groups = {
|
||||
# "policy": ["policy", "images"],
|
||||
# "critic": ["policy", "privileged"],
|
||||
# }
|
||||
|
||||
# This way, the policy will receive the "policy" and "images" observations, and the critic will
|
||||
# receive the "policy" and "privileged" observations.
|
||||
|
||||
# For more details, please check ``vec_env.py`` in the rsl_rl library.
|
||||
# """
|
||||
|
||||
# clip_actions: float | None = None
|
||||
# """The clipping value for actions. If None, then no clipping is done. Defaults to None.
|
||||
|
||||
# .. note::
|
||||
# This clipping is performed inside the :class:`RslRlVecEnvWrapper` wrapper.
|
||||
# """
|
||||
|
||||
save_interval: int = MISSING
|
||||
"""The number of iterations between saves."""
|
||||
|
||||
experiment_name: str = MISSING
|
||||
"""The experiment name."""
|
||||
|
||||
run_name: str = ""
|
||||
"""The run name. Default is empty string.
|
||||
|
||||
The name of the run directory is typically the time-stamp at execution. If the run name is not empty,
|
||||
then it is appended to the run directory's name, i.e. the logging directory's name will become
|
||||
``{time-stamp}_{run_name}``.
|
||||
"""
|
||||
|
||||
logger: Literal["tensorboard", "neptune", "wandb"] = "tensorboard"
|
||||
"""The logger to use. Default is tensorboard."""
|
||||
|
||||
neptune_project: str = "isaaclab"
|
||||
"""The neptune project name. Default is "isaaclab"."""
|
||||
|
||||
wandb_project: str = "isaaclab"
|
||||
"""The wandb project name. Default is "isaaclab"."""
|
||||
|
||||
resume: bool = False
|
||||
"""Whether to resume a previous training. Default is False.
|
||||
|
||||
This flag will be ignored for distillation.
|
||||
"""
|
||||
|
||||
load_run: str = ".*"
|
||||
"""The run directory to load. Default is ".*" (all).
|
||||
|
||||
If regex expression, the latest (alphabetical order) matching run will be loaded.
|
||||
"""
|
||||
|
||||
load_checkpoint: str = "model_.*.pt"
|
||||
"""The checkpoint file to load. Default is ``"model_.*.pt"`` (all).
|
||||
|
||||
If regex expression, the latest (alphabetical order) matching file will be loaded.
|
||||
"""
|
||||
|
||||
|
||||
@configclass
|
||||
class HIMPPOActorCriticCfg:
|
||||
"""Configuration of the HIM PPO actor-critic."""
|
||||
|
||||
actor_hidden_dims: list[int] = [512, 256, 128]
|
||||
"""The hidden dimensions of the actor network."""
|
||||
|
||||
critic_hidden_dims: list[int] = [512, 256, 128]
|
||||
"""The hidden dimensions of the critic network."""
|
||||
|
||||
activation: str = "elu"
|
||||
"""The activation function to use. Default is 'elu'."""
|
||||
|
||||
init_noise_std: float = 1.0
|
||||
"""The initial noise standard deviation for the actor. Default is 1.0."""
|
||||
|
||||
normalize_obs: bool = False
|
||||
"""Whether to normalize observations. Default is False."""
|
||||
|
||||
@configclass
|
||||
class HIMPPPOAlgorithmCfg:
|
||||
"""Configuration of the HIM PPO algorithm."""
|
||||
num_learning_epochs: int = 1
|
||||
"""The number of learning epochs per update. Default is 1."""
|
||||
|
||||
num_mini_batches: int = 1
|
||||
"""The number of mini-batches per update. Default is 1."""
|
||||
|
||||
clip_param: float = 0.2
|
||||
"""The clipping parameter for PPO. Default is 0.2."""
|
||||
|
||||
gamma: float = 0.998
|
||||
"""The discount factor. Default is 0.998."""
|
||||
|
||||
lam: float = 0.95
|
||||
"""The GAE lambda parameter. Default is 0.95."""
|
||||
|
||||
value_loss_coef: float = 1.0
|
||||
"""The coefficient for the value loss. Default is 1.0."""
|
||||
|
||||
entropy_coef: float = 0.0
|
||||
"""The coefficient for the entropy bonus. Default is 0.0."""
|
||||
|
||||
learning_rate: float = 1.0e-3
|
||||
"""The learning rate. Default is 1.0e-3."""
|
||||
|
||||
max_grad_norm: float = 1.0
|
||||
"""The maximum gradient norm for clipping. Default is 1.0."""
|
||||
|
||||
use_clipped_value_loss: bool = True
|
||||
"""Whether to use clipped value loss. Default is True."""
|
||||
|
||||
schedule: str = "fixed"
|
||||
"""The learning rate schedule. Default is 'fixed'."""
|
||||
|
||||
desired_kl: float = 0.01
|
||||
"""The desired KL divergence for adaptive learning rate. Default is 0.01."""
|
||||
|
||||
@configclass
|
||||
class HIMOnPolicyRunnerCfg(HIMBaseRunnerCfg):
|
||||
"""Configuration of the runner for on-policy algorithms."""
|
||||
|
||||
class_name: str = "HIMOnPolicyRunner"
|
||||
"""The runner class name. Default is OnPolicyRunner."""
|
||||
|
||||
policy_class_name: str = "HIMActorCritic"
|
||||
"""The policy class name. Default is HIMActorCritic."""
|
||||
|
||||
algorithm_class_name: str = "HIMPPO"
|
||||
"""The algorithm class name. Default is HIMPPO."""
|
||||
|
||||
policy: HIMPPOActorCriticCfg = MISSING
|
||||
"""The policy configuration."""
|
||||
|
||||
algorithm: HIMPPPOAlgorithmCfg = MISSING
|
||||
"""The algorithm configuration."""
|
||||
|
||||
history_length: int = 0
|
||||
"""Number of historical time steps to stack with current observation (0 means current only). Default is 0."""
|
||||
|
||||
privileged_history_length: int = 0
|
||||
"""Number of historical time steps to stack with current privileged observation. Default is 0."""
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .vec_env import VecEnv
|
||||
@@ -0,0 +1,60 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import torch
|
||||
from typing import Tuple, Union
|
||||
|
||||
# minimal interface of the environment
|
||||
class VecEnv(ABC):
|
||||
num_envs: int
|
||||
num_obs: int
|
||||
num_privileged_obs: int
|
||||
num_actions: int
|
||||
max_episode_length: int
|
||||
privileged_obs_buf: torch.Tensor
|
||||
obs_buf: torch.Tensor
|
||||
rew_buf: torch.Tensor
|
||||
reset_buf: torch.Tensor
|
||||
episode_length_buf: torch.Tensor # current episode duration
|
||||
extras: dict
|
||||
device: torch.device
|
||||
@abstractmethod
|
||||
def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, Union[torch.Tensor, None], torch.Tensor, torch.Tensor, dict]:
|
||||
pass
|
||||
@abstractmethod
|
||||
def reset(self, env_ids: Union[list, torch.Tensor]):
|
||||
pass
|
||||
@abstractmethod
|
||||
def get_observations(self) -> torch.Tensor:
|
||||
pass
|
||||
@abstractmethod
|
||||
def get_privileged_observations(self) -> Union[torch.Tensor, None]:
|
||||
pass
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .him_actor_critic import HIMActorCritic
|
||||
from .him_estimator import HIMEstimator
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
from ..modules.him_estimator import HIMEstimator
|
||||
|
||||
class RunningMeanStd:
|
||||
# Dynamically calculate mean and std
|
||||
def __init__(self, shape, device): # shape:the dimension of input data
|
||||
self.n = 1e-4
|
||||
self.uninitialized = True
|
||||
self.mean = torch.zeros(shape, device=device)
|
||||
self.var = torch.ones(shape, device=device)
|
||||
|
||||
def update(self, x):
|
||||
count = self.n
|
||||
batch_count = x.size(0)
|
||||
tot_count = count + batch_count
|
||||
|
||||
old_mean = self.mean.clone()
|
||||
delta = torch.mean(x, dim=0) - old_mean
|
||||
|
||||
self.mean = old_mean + delta * batch_count / tot_count
|
||||
m_a = self.var * count
|
||||
m_b = x.var(dim=0) * batch_count
|
||||
M2 = m_a + m_b + torch.square(delta) * count * batch_count / tot_count
|
||||
self.var = M2 / tot_count
|
||||
self.n = tot_count
|
||||
|
||||
class Normalization:
|
||||
def __init__(self, shape, device='cuda:0'):
|
||||
self.running_ms = RunningMeanStd(shape=shape, device=device)
|
||||
|
||||
def __call__(self, x, update=False):
|
||||
# Whether to update the mean and std,during the evaluating,update=Flase
|
||||
if update:
|
||||
self.running_ms.update(x)
|
||||
x = (x - self.running_ms.mean) / (torch.sqrt(self.running_ms.var) + 1e-4)
|
||||
|
||||
return x
|
||||
|
||||
class HIMActorCritic(nn.Module):
|
||||
is_recurrent = False
|
||||
def __init__(self, num_actor_obs,
|
||||
num_critic_obs,
|
||||
num_one_step_obs,
|
||||
num_actions,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation='elu',
|
||||
init_noise_std=1.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
|
||||
super(HIMActorCritic, self).__init__()
|
||||
|
||||
activation = get_activation(activation)
|
||||
|
||||
self.history_size = int(num_actor_obs/num_one_step_obs)
|
||||
self.num_actor_obs = num_actor_obs
|
||||
self.num_actions = num_actions
|
||||
self.num_one_step_obs = num_one_step_obs
|
||||
|
||||
mlp_input_dim_a = num_one_step_obs + 3 + 16
|
||||
mlp_input_dim_c = num_critic_obs
|
||||
|
||||
# Estimator
|
||||
self.estimator = HIMEstimator(temporal_steps=self.history_size, num_one_step_obs=num_one_step_obs)
|
||||
|
||||
# Policy
|
||||
actor_layers = []
|
||||
actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
|
||||
actor_layers.append(activation)
|
||||
for l in range(len(actor_hidden_dims)):
|
||||
if l == len(actor_hidden_dims) - 1:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
|
||||
# actor_layers.append(nn.Tanh())
|
||||
else:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
|
||||
actor_layers.append(activation)
|
||||
self.actor = nn.Sequential(*actor_layers)
|
||||
|
||||
# Value function
|
||||
critic_layers = []
|
||||
critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
|
||||
critic_layers.append(activation)
|
||||
for l in range(len(critic_hidden_dims)):
|
||||
if l == len(critic_hidden_dims) - 1:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
|
||||
else:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
|
||||
critic_layers.append(activation)
|
||||
self.critic = nn.Sequential(*critic_layers)
|
||||
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
print(f'Estimator: {self.estimator.encoder}')
|
||||
|
||||
# Action noise
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
# disable args validation for speedup
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
# seems that we get better performance without init
|
||||
# self.init_memory_weights(self.memory_a, 0.001, 0.)
|
||||
# self.init_memory_weights(self.memory_c, 0.001, 0.)
|
||||
|
||||
@staticmethod
|
||||
# not used at the moment
|
||||
def init_weights(sequential, scales):
|
||||
[torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in
|
||||
enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]
|
||||
|
||||
|
||||
def reset(self, dones=None):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, obs_history):
|
||||
with torch.no_grad():
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actor_input = torch.cat((obs_history[:,:self.num_one_step_obs], vel, latent), dim=-1)
|
||||
|
||||
# Check for NaN/Inf in inputs
|
||||
if torch.isnan(actor_input).any() or torch.isinf(actor_input).any():
|
||||
print(f"[ERROR] NaN/Inf detected in actor_input before normalization!")
|
||||
print(f" - obs_history stats: min={obs_history.min():.4f}, max={obs_history.max():.4f}, mean={obs_history.mean():.4f}, has_nan={torch.isnan(obs_history).any()}")
|
||||
print(f" - vel stats: min={vel.min():.4f}, max={vel.max():.4f}, mean={vel.mean():.4f}, has_nan={torch.isnan(vel).any()}")
|
||||
print(f" - latent stats: min={latent.min():.4f}, max={latent.max():.4f}, mean={latent.mean():.4f}, has_nan={torch.isnan(latent).any()}")
|
||||
raise ValueError("NaN/Inf in actor_input before normalization")
|
||||
|
||||
mean = self.actor(actor_input)
|
||||
|
||||
# Check for NaN/Inf in actor output
|
||||
if torch.isnan(mean).any() or torch.isinf(mean).any():
|
||||
print(f"[ERROR] NaN/Inf in actor output (mean)!")
|
||||
print(f" - actor_input stats: min={actor_input.min():.4f}, max={actor_input.max():.4f}, mean={actor_input.mean():.4f}")
|
||||
print(f" - Batch size: {actor_input.shape[0]}")
|
||||
|
||||
# Check actor network weights AFTER forward (they should still be OK)
|
||||
for name, param in self.actor.named_parameters():
|
||||
if torch.isnan(param).any():
|
||||
print(f" - NaN in actor weight: {name}")
|
||||
if torch.isinf(param).any():
|
||||
print(f" - Inf in actor weight: {name}")
|
||||
|
||||
raise ValueError("NaN or Inf detected in actor network output!")
|
||||
|
||||
self.distribution = Normal(mean, mean*0. + self.std)
|
||||
|
||||
def act(self, obs_history=None, **kwargs):
|
||||
self.update_distribution(obs_history)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, obs_history, observations=None):
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actions_mean = self.actor(torch.cat((obs_history[:,:self.num_one_step_obs], vel, latent), dim=-1))
|
||||
return actions_mean
|
||||
|
||||
def test_inference(self, obs_history, observations=None):
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actions_mean = self.actor(torch.cat((obs_history[:,:self.num_one_step_obs], vel, latent), dim=-1))
|
||||
estimator_output = torch.cat((vel, latent), dim=-1)
|
||||
return actions_mean, estimator_output
|
||||
|
||||
def evaluate(self, critic_observations, **kwargs):
|
||||
value = self.critic(critic_observations)
|
||||
return value
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
@@ -0,0 +1,155 @@
|
||||
import copy
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import torch.nn.functional as F
|
||||
import torch.distributions as torchd
|
||||
from torch.distributions import Normal, Categorical
|
||||
|
||||
|
||||
class HIMEstimator(nn.Module):
|
||||
def __init__(self,
|
||||
temporal_steps,
|
||||
num_one_step_obs,
|
||||
enc_hidden_dims=[128, 64, 16],
|
||||
tar_hidden_dims=[128, 64],
|
||||
activation='elu',
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=10.0,
|
||||
num_prototype=32,
|
||||
temperature=3.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("Estimator_CL.__init__ got unexpected arguments, which will be ignored: " + str(
|
||||
[key for key in kwargs.keys()]))
|
||||
super(HIMEstimator, self).__init__()
|
||||
activation = get_activation(activation)
|
||||
|
||||
self.temporal_steps = temporal_steps
|
||||
self.num_one_step_obs = num_one_step_obs
|
||||
self.num_latent = enc_hidden_dims[-1]
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.temperature = temperature
|
||||
|
||||
# Encoder
|
||||
enc_input_dim = self.temporal_steps * self.num_one_step_obs
|
||||
enc_layers = []
|
||||
for l in range(len(enc_hidden_dims) - 1):
|
||||
enc_layers += [nn.Linear(enc_input_dim, enc_hidden_dims[l]), activation]
|
||||
enc_input_dim = enc_hidden_dims[l]
|
||||
enc_layers += [nn.Linear(enc_input_dim, enc_hidden_dims[-1] + 3)]
|
||||
self.encoder = nn.Sequential(*enc_layers)
|
||||
|
||||
# Target
|
||||
tar_input_dim = self.num_one_step_obs
|
||||
tar_layers = []
|
||||
for l in range(len(tar_hidden_dims)):
|
||||
tar_layers += [nn.Linear(tar_input_dim, tar_hidden_dims[l]), activation]
|
||||
tar_input_dim = tar_hidden_dims[l]
|
||||
tar_layers += [nn.Linear(tar_input_dim, enc_hidden_dims[-1])]
|
||||
self.target = nn.Sequential(*tar_layers)
|
||||
|
||||
# Prototype
|
||||
self.proto = nn.Embedding(num_prototype, enc_hidden_dims[-1])
|
||||
|
||||
# Optimizer
|
||||
self.learning_rate = learning_rate
|
||||
self.optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
|
||||
|
||||
def get_latent(self, obs_history):
|
||||
vel, z = self.encode(obs_history)
|
||||
return vel.detach(), z.detach()
|
||||
|
||||
def forward(self, obs_history):
|
||||
parts = self.encoder(obs_history.detach())
|
||||
vel, z = parts[..., :3], parts[..., 3:]
|
||||
z = F.normalize(z, dim=-1, p=2)
|
||||
return vel.detach(), z.detach()
|
||||
|
||||
def encode(self, obs_history):
|
||||
parts = self.encoder(obs_history.detach())
|
||||
vel, z = parts[..., :3], parts[..., 3:]
|
||||
z = F.normalize(z, dim=-1, p=2)
|
||||
return vel, z
|
||||
|
||||
def update(self, obs_history, next_critic_obs, lr=None):
|
||||
if lr is not None:
|
||||
self.learning_rate = lr
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
vel = next_critic_obs[:, self.num_one_step_obs:self.num_one_step_obs+3].detach()
|
||||
next_obs = next_critic_obs.detach()[:, 3:self.num_one_step_obs+3]
|
||||
|
||||
z_s = self.encoder(obs_history)
|
||||
z_t = self.target(next_obs)
|
||||
pred_vel, z_s = z_s[..., :3], z_s[..., 3:]
|
||||
|
||||
z_s = F.normalize(z_s, dim=-1, p=2)
|
||||
z_t = F.normalize(z_t, dim=-1, p=2)
|
||||
|
||||
with torch.no_grad():
|
||||
w = self.proto.weight.data.clone()
|
||||
w = F.normalize(w, dim=-1, p=2)
|
||||
self.proto.weight.copy_(w)
|
||||
|
||||
score_s = z_s @ self.proto.weight.T
|
||||
score_t = z_t @ self.proto.weight.T
|
||||
|
||||
with torch.no_grad():
|
||||
q_s = sinkhorn(score_s)
|
||||
q_t = sinkhorn(score_t)
|
||||
|
||||
log_p_s = F.log_softmax(score_s / self.temperature, dim=-1)
|
||||
log_p_t = F.log_softmax(score_t / self.temperature, dim=-1)
|
||||
|
||||
swap_loss = -0.5 * (q_s * log_p_t + q_t * log_p_s).mean()
|
||||
estimation_loss = F.mse_loss(pred_vel, vel)
|
||||
losses = estimation_loss + swap_loss
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
losses.backward()
|
||||
nn.utils.clip_grad_norm_(self.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
return estimation_loss.item(), swap_loss.item()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sinkhorn(out, eps=0.05, iters=3):
|
||||
Q = torch.exp(out / eps).T
|
||||
K, B = Q.shape[0], Q.shape[1]
|
||||
Q /= Q.sum()
|
||||
|
||||
for it in range(iters):
|
||||
# normalize each row: total weight per prototype must be 1/K
|
||||
Q /= torch.sum(Q, dim=1, keepdim=True)
|
||||
Q /= K
|
||||
|
||||
# normalize each column: total weight per sample must be 1/B
|
||||
Q /= torch.sum(Q, dim=0, keepdim=True)
|
||||
Q /= B
|
||||
return (Q * B).T
|
||||
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "silu":
|
||||
return nn.SiLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .him_on_policy_runner import HIMOnPolicyRunner
|
||||
@@ -0,0 +1,362 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import time
|
||||
import os
|
||||
from collections import deque
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
|
||||
from ..algorithms import HIMPPO
|
||||
from ..modules import HIMActorCritic
|
||||
from ..env import VecEnv
|
||||
|
||||
|
||||
class HIMOnPolicyRunner:
|
||||
|
||||
def __init__(self,
|
||||
env: VecEnv,
|
||||
train_cfg,
|
||||
log_dir=None,
|
||||
device='cpu'):
|
||||
"""Initialize HimLoco on-policy runner.
|
||||
|
||||
Args:
|
||||
env: Vectorized environment.
|
||||
train_cfg: Training configuration dictionary.
|
||||
log_dir: Directory for logging.
|
||||
device: Device to run on.
|
||||
"""
|
||||
# Store configuration
|
||||
self.cfg = train_cfg # Full training config (for logging compatibility)
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
|
||||
# Determine observation dimensions
|
||||
if self.env.num_privileged_obs is not None:
|
||||
num_critic_obs = self.env.num_privileged_obs
|
||||
else:
|
||||
num_critic_obs = self.env.num_obs
|
||||
self.num_actor_obs = self.env.num_obs
|
||||
self.num_critic_obs = num_critic_obs
|
||||
|
||||
# Initialize policy network
|
||||
actor_critic_class = eval(train_cfg["policy_class_name"]) # HIMActorCritic
|
||||
actor_critic: HIMActorCritic = actor_critic_class(
|
||||
self.env.num_obs, # historical obs
|
||||
num_critic_obs, # historical privileged obs
|
||||
self.env.num_one_step_obs,
|
||||
self.env.num_actions,
|
||||
**self.policy_cfg
|
||||
).to(self.device)
|
||||
|
||||
# Initialize algorithm
|
||||
alg_class = eval(train_cfg["algorithm_class_name"]) # HIMPPO
|
||||
self.alg: HIMPPO = alg_class(actor_critic, device=self.device, **self.alg_cfg)
|
||||
|
||||
# Training configuration
|
||||
self.num_steps_per_env = train_cfg["num_steps_per_env"]
|
||||
self.save_interval = train_cfg["save_interval"]
|
||||
|
||||
# Initialize storage and model
|
||||
self.alg.init_storage(
|
||||
self.env.num_envs,
|
||||
self.num_steps_per_env,
|
||||
[self.env.num_obs],
|
||||
[self.env.num_privileged_obs],
|
||||
[self.env.num_actions]
|
||||
)
|
||||
|
||||
# Logging
|
||||
self.log_dir = log_dir
|
||||
self.writer = None
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.current_learning_iteration = 0
|
||||
self.logger_type = self.cfg["logger"].lower()
|
||||
|
||||
# _, _ = self.env.reset() we call this in wrapper
|
||||
|
||||
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
|
||||
"""Train the policy using HimLoco PPO algorithm.
|
||||
|
||||
Args:
|
||||
num_learning_iterations: Number of policy updates.
|
||||
init_at_random_ep_len: Whether to randomize initial episode lengths.
|
||||
"""
|
||||
# Initialize writer (support multiple logger types like RSL-RL)
|
||||
if self.log_dir is not None and self.writer is None:
|
||||
# Launch either Tensorboard, WandB, or Neptune summary writer(s)
|
||||
if self.logger_type == "neptune":
|
||||
from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter
|
||||
self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.alg_cfg)
|
||||
# TODO: Add config logging support
|
||||
# self.writer.log_config(self.env.cfg, train_cfg, self.alg_cfg, self.policy_cfg)
|
||||
elif self.logger_type == "wandb":
|
||||
from rsl_rl.utils.wandb_utils import WandbSummaryWriter
|
||||
self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.alg_cfg)
|
||||
# TODO: Add config logging support
|
||||
elif self.logger_type == "tensorboard":
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
|
||||
else:
|
||||
raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.")
|
||||
|
||||
if init_at_random_ep_len:
|
||||
self.env.episode_length_buf = torch.randint_like(self.env.episode_length_buf, high=int(self.env.max_episode_length))
|
||||
obs = self.env.get_observations()
|
||||
privileged_obs = self.env.get_privileged_observations()
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs = obs.to(self.device), critic_obs.to(self.device)
|
||||
self.alg.actor_critic.train() # switch to train mode (for dropout for example)
|
||||
|
||||
ep_infos = []
|
||||
rewbuffer = deque(maxlen=100)
|
||||
lenbuffer = deque(maxlen=100)
|
||||
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
start_iter = self.current_learning_iteration
|
||||
tot_iter = start_iter + num_learning_iterations
|
||||
for it in range(start_iter, tot_iter):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for i in range(self.num_steps_per_env):
|
||||
actions = self.alg.act(obs, critic_obs)
|
||||
obs, privileged_obs, rewards, dones, infos, termination_ids, termination_privileged_obs = self.env.step(actions)
|
||||
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs, rewards, dones = obs.to(self.device), critic_obs.to(self.device), rewards.to(self.device), dones.to(self.device)
|
||||
termination_ids = termination_ids.to(self.device)
|
||||
termination_privileged_obs = termination_privileged_obs.to(self.device)
|
||||
|
||||
next_critic_obs = critic_obs.clone().detach()
|
||||
next_critic_obs[termination_ids] = termination_privileged_obs.clone().detach()
|
||||
|
||||
self.alg.process_env_step(rewards, dones, infos, next_critic_obs)
|
||||
# print("+++++++++++++++++++++++++")
|
||||
# print(obs[0])
|
||||
# print(next_critic_obs[0, :self.env.num_one_step_obs])
|
||||
# print("+++++++++++++++++++++++++")
|
||||
if self.log_dir is not None:
|
||||
# Book keeping
|
||||
if 'episode' in infos:
|
||||
ep_infos.append(infos['episode'])
|
||||
elif 'log' in infos:
|
||||
ep_infos.append(infos['log'])
|
||||
# Update rewards
|
||||
cur_reward_sum += rewards
|
||||
# Update episode length
|
||||
cur_episode_length += 1
|
||||
# Clear data for completed episodes
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
|
||||
cur_reward_sum[new_ids] = 0
|
||||
cur_episode_length[new_ids] = 0
|
||||
|
||||
stop = time.time()
|
||||
collection_time = stop - start
|
||||
|
||||
# Learning step
|
||||
start = stop
|
||||
self.alg.compute_returns(critic_obs)
|
||||
|
||||
mean_value_loss, mean_surrogate_loss, mean_estimation_loss, mean_swap_loss = self.alg.update()
|
||||
stop = time.time()
|
||||
learn_time = stop - start
|
||||
self.current_learning_iteration = it
|
||||
# log info
|
||||
if self.log_dir is not None:
|
||||
self.log(locals())
|
||||
# Save model
|
||||
if it % self.save_interval == 0:
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(it)))
|
||||
# Clear episode infos
|
||||
ep_infos.clear()
|
||||
|
||||
# Save the final model after training
|
||||
if self.log_dir is not None:
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)))
|
||||
|
||||
def log(self, locs, width=80, pad=35):
|
||||
"""Log training information to console and TensorBoard.
|
||||
|
||||
Enhanced version compatible with RSL-RL style logging, supporting:
|
||||
- Dynamic episode info processing
|
||||
- Loss dictionary iteration
|
||||
- Detailed performance metrics
|
||||
"""
|
||||
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
|
||||
self.tot_time += locs['collection_time'] + locs['learn_time']
|
||||
iteration_time = locs['collection_time'] + locs['learn_time']
|
||||
|
||||
# -- Episode info
|
||||
ep_string = ""
|
||||
if locs['ep_infos']:
|
||||
for key in locs['ep_infos'][0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
for ep_info in locs['ep_infos']:
|
||||
# skip missing keys
|
||||
if key not in ep_info:
|
||||
continue
|
||||
# handle scalar and zero dimensional tensor infos
|
||||
if not isinstance(ep_info[key], torch.Tensor):
|
||||
ep_info[key] = torch.Tensor([ep_info[key]])
|
||||
if len(ep_info[key].shape) == 0:
|
||||
ep_info[key] = ep_info[key].unsqueeze(0)
|
||||
infotensor = torch.cat((infotensor, ep_info[key].to(self.device)))
|
||||
|
||||
if infotensor.numel() > 0:
|
||||
value = torch.mean(infotensor)
|
||||
# log to tensorboard with proper namespace
|
||||
if "/" in key:
|
||||
self.writer.add_scalar(key, value, locs['it'])
|
||||
ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n"""
|
||||
else:
|
||||
self.writer.add_scalar('Episode/' + key, value, locs['it'])
|
||||
ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
|
||||
|
||||
# -- Policy metrics
|
||||
mean_std = self.alg.actor_critic.std.mean()
|
||||
fps = int(self.num_steps_per_env * self.env.num_envs / (locs['collection_time'] + locs['learn_time']))
|
||||
|
||||
# -- Losses
|
||||
loss_dict = {
|
||||
'value_function': locs['mean_value_loss'],
|
||||
'surrogate': locs['mean_surrogate_loss'],
|
||||
'estimation': locs['mean_estimation_loss'],
|
||||
'swap': locs['mean_swap_loss'],
|
||||
}
|
||||
for key, value in loss_dict.items():
|
||||
self.writer.add_scalar(f'Loss/{key}', value, locs['it'])
|
||||
|
||||
self.writer.add_scalar('Loss/learning_rate', self.alg.learning_rate, locs['it'])
|
||||
self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), locs['it'])
|
||||
|
||||
# -- Performance
|
||||
self.writer.add_scalar('Perf/total_fps', fps, locs['it'])
|
||||
self.writer.add_scalar('Perf/collection_time', locs['collection_time'], locs['it'])
|
||||
self.writer.add_scalar('Perf/learning_time', locs['learn_time'], locs['it'])
|
||||
|
||||
# -- Training metrics
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
self.writer.add_scalar('Train/mean_reward', statistics.mean(locs['rewbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_episode_length', statistics.mean(locs['lenbuffer']), locs['it'])
|
||||
if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging
|
||||
self.writer.add_scalar('Train/mean_reward/time', statistics.mean(locs['rewbuffer']), self.tot_time)
|
||||
self.writer.add_scalar('Train/mean_episode_length/time', statistics.mean(locs['lenbuffer']), self.tot_time)
|
||||
|
||||
# -- Console output
|
||||
str_header = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m "
|
||||
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
log_string = (
|
||||
f"""{'#' * width}\n"""
|
||||
f"""{str_header.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs['collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
)
|
||||
# Add losses
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n"""
|
||||
# Add rewards
|
||||
log_string += (
|
||||
f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n"""
|
||||
f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n"""
|
||||
)
|
||||
else:
|
||||
log_string = (
|
||||
f"""{'#' * width}\n"""
|
||||
f"""{str_header.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs['collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
)
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n"""
|
||||
|
||||
# Add episode info
|
||||
log_string += ep_string
|
||||
|
||||
# Add summary
|
||||
log_string += (
|
||||
f"""{'-' * width}\n"""
|
||||
f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n"""
|
||||
f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n"""
|
||||
f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n"""
|
||||
f"""{'ETA:':>{pad}} {time.strftime(
|
||||
"%H:%M:%S",
|
||||
time.gmtime(
|
||||
self.tot_time / (locs['it'] - locs['start_iter'] + 1)
|
||||
* (locs['start_iter'] + locs['num_learning_iterations'] - locs['it'])
|
||||
)
|
||||
)}\n"""
|
||||
)
|
||||
print(log_string)
|
||||
|
||||
def save(self, path, infos=None):
|
||||
"""Save model checkpoint.
|
||||
|
||||
Args:
|
||||
path: Path to save the model.
|
||||
infos: Additional information to save.
|
||||
"""
|
||||
saved_dict = {
|
||||
'model_state_dict': self.alg.actor_critic.state_dict(),
|
||||
'optimizer_state_dict': self.alg.optimizer.state_dict(),
|
||||
'estimator_optimizer_state_dict': self.alg.actor_critic.estimator.optimizer.state_dict(),
|
||||
'iter': self.current_learning_iteration,
|
||||
'infos': infos,
|
||||
}
|
||||
torch.save(saved_dict, path)
|
||||
|
||||
# Upload model to external logging service
|
||||
if hasattr(self, 'logger_type') and self.logger_type in ["neptune", "wandb"]:
|
||||
if hasattr(self.writer, 'save_model'):
|
||||
self.writer.save_model(path, self.current_learning_iteration)
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
loaded_dict = torch.load(path)
|
||||
self.alg.actor_critic.load_state_dict(loaded_dict['model_state_dict'])
|
||||
if load_optimizer:
|
||||
self.alg.optimizer.load_state_dict(loaded_dict['optimizer_state_dict'])
|
||||
self.alg.actor_critic.estimator.optimizer.load_state_dict(loaded_dict['estimator_optimizer_state_dict'])
|
||||
self.current_learning_iteration = loaded_dict['iter']
|
||||
return loaded_dict['infos']
|
||||
|
||||
def get_inference_policy(self, device=None):
|
||||
self.alg.actor_critic.eval() # switch to evaluation mode (dropout for example)
|
||||
if device is not None:
|
||||
self.alg.actor_critic.to(device)
|
||||
return self.alg.actor_critic.act_inference
|
||||
@@ -0,0 +1,4 @@
|
||||
# Copyright 2021 ETH Zurich, NVIDIA CORPORATION
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from .him_rollout_storage import HIMRolloutStorage
|
||||
@@ -0,0 +1,167 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
# from ..utils import split_and_pad_trajectories
|
||||
|
||||
class HIMRolloutStorage:
|
||||
class Transition:
|
||||
def __init__(self):
|
||||
self.observations = None
|
||||
self.critic_observations = None
|
||||
self.actions = None
|
||||
self.rewards = None
|
||||
self.dones = None
|
||||
self.values = None
|
||||
self.actions_log_prob = None
|
||||
self.action_mean = None
|
||||
self.action_sigma = None
|
||||
self.next_critic_observations = None
|
||||
|
||||
def clear(self):
|
||||
self.__init__()
|
||||
|
||||
def __init__(self, num_envs, num_transitions_per_env, obs_shape, privileged_obs_shape, actions_shape, device='cpu'):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.obs_shape = obs_shape
|
||||
self.privileged_obs_shape = privileged_obs_shape
|
||||
self.actions_shape = actions_shape
|
||||
|
||||
# Core
|
||||
self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device)
|
||||
if privileged_obs_shape[0] is not None:
|
||||
self.privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
self.next_privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
else:
|
||||
self.privileged_observations = None
|
||||
self.next_privileged_observations = None
|
||||
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
|
||||
|
||||
# For PPO
|
||||
self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
|
||||
self.num_transitions_per_env = num_transitions_per_env
|
||||
self.num_envs = num_envs
|
||||
|
||||
self.step = 0
|
||||
|
||||
def add_transitions(self, transition: Transition):
|
||||
if self.step >= self.num_transitions_per_env:
|
||||
raise AssertionError("Rollout buffer overflow")
|
||||
self.observations[self.step].copy_(transition.observations)
|
||||
if self.privileged_observations is not None: self.privileged_observations[self.step].copy_(transition.critic_observations)
|
||||
if self.next_privileged_observations is not None: self.next_privileged_observations[self.step].copy_(transition.next_critic_observations)
|
||||
self.actions[self.step].copy_(transition.actions)
|
||||
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
|
||||
self.dones[self.step].copy_(transition.dones.view(-1, 1))
|
||||
self.values[self.step].copy_(transition.values)
|
||||
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
|
||||
self.mu[self.step].copy_(transition.action_mean)
|
||||
self.sigma[self.step].copy_(transition.action_sigma)
|
||||
self.step += 1
|
||||
|
||||
def clear(self):
|
||||
self.step = 0
|
||||
|
||||
def compute_returns(self, last_values, gamma, lam):
|
||||
advantage = 0
|
||||
for step in reversed(range(self.num_transitions_per_env)):
|
||||
if step == self.num_transitions_per_env - 1:
|
||||
next_values = last_values
|
||||
else:
|
||||
next_values = self.values[step + 1]
|
||||
next_is_not_terminal = 1.0 - self.dones[step].float()
|
||||
delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
|
||||
advantage = delta + next_is_not_terminal * gamma * lam * advantage
|
||||
self.returns[step] = advantage + self.values[step]
|
||||
|
||||
# Compute and normalize the advantages
|
||||
self.advantages = self.returns - self.values
|
||||
self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
|
||||
|
||||
def get_statistics(self):
|
||||
done = self.dones
|
||||
done[-1] = 1
|
||||
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]))
|
||||
trajectory_lengths = (done_indices[1:] - done_indices[:-1])
|
||||
return trajectory_lengths.float().mean(), self.rewards.mean()
|
||||
|
||||
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
batch_size = self.num_envs * self.num_transitions_per_env
|
||||
mini_batch_size = batch_size // num_mini_batches
|
||||
indices = torch.randperm(num_mini_batches*mini_batch_size, requires_grad=False, device=self.device)
|
||||
|
||||
observations = self.observations.flatten(0, 1)
|
||||
if self.privileged_observations is not None:
|
||||
critic_observations = self.privileged_observations.flatten(0, 1)
|
||||
next_critic_observations = self.next_privileged_observations.flatten(0, 1)
|
||||
else:
|
||||
critic_observations = observations
|
||||
next_critic_observations = observations
|
||||
|
||||
actions = self.actions.flatten(0, 1)
|
||||
values = self.values.flatten(0, 1)
|
||||
returns = self.returns.flatten(0, 1)
|
||||
old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
|
||||
advantages = self.advantages.flatten(0, 1)
|
||||
old_mu = self.mu.flatten(0, 1)
|
||||
old_sigma = self.sigma.flatten(0, 1)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
|
||||
start = i*mini_batch_size
|
||||
end = (i+1)*mini_batch_size
|
||||
batch_idx = indices[start:end]
|
||||
|
||||
obs_batch = observations[batch_idx]
|
||||
next_critic_observations_batch = next_critic_observations[batch_idx]
|
||||
critic_observations_batch = critic_observations[batch_idx]
|
||||
actions_batch = actions[batch_idx]
|
||||
target_values_batch = values[batch_idx]
|
||||
returns_batch = returns[batch_idx]
|
||||
old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
|
||||
advantages_batch = advantages[batch_idx]
|
||||
old_mu_batch = old_mu[batch_idx]
|
||||
old_sigma_batch = old_sigma[batch_idx]
|
||||
yield obs_batch, critic_observations_batch, actions_batch, next_critic_observations_batch, target_values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .utils import split_and_pad_trajectories, unpad_trajectories
|
||||
@@ -0,0 +1,71 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
|
||||
def split_and_pad_trajectories(tensor, dones):
|
||||
""" Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory.
|
||||
Returns masks corresponding to valid parts of the trajectories
|
||||
Example:
|
||||
Input: [ [a1, a2, a3, a4 | a5, a6],
|
||||
[b1, b2 | b3, b4, b5 | b6]
|
||||
]
|
||||
|
||||
Output:[ [a1, a2, a3, a4], | [ [True, True, True, True],
|
||||
[a5, a6, 0, 0], | [True, True, False, False],
|
||||
[b1, b2, 0, 0], | [True, True, False, False],
|
||||
[b3, b4, b5, 0], | [True, True, True, False],
|
||||
[b6, 0, 0, 0] | [True, False, False, False],
|
||||
] | ]
|
||||
|
||||
Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions]
|
||||
"""
|
||||
dones = dones.clone()
|
||||
dones[-1] = 1
|
||||
# Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping
|
||||
flat_dones = dones.transpose(1, 0).reshape(-1, 1)
|
||||
|
||||
# Get length of trajectory by counting the number of successive not done elements
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
|
||||
trajectory_lengths = done_indices[1:] - done_indices[:-1]
|
||||
trajectory_lengths_list = trajectory_lengths.tolist()
|
||||
# Extract the individual trajectories
|
||||
trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list)
|
||||
padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)
|
||||
|
||||
|
||||
trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1)
|
||||
return padded_trajectories, trajectory_masks
|
||||
|
||||
def unpad_trajectories(trajectories, masks):
|
||||
""" Does the inverse operation of split_and_pad_trajectories()
|
||||
"""
|
||||
# Need to transpose before and after the masking to have proper reshaping
|
||||
return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0)
|
||||
@@ -0,0 +1,113 @@
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg, BuiltinVelocityActuatorCfg
|
||||
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.utils.spec_config import CollisionCfg
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).parent.parent.parent
|
||||
ROBOT_XML = _PROJECT_ROOT / "mjcf" / "wheelleg.xml"
|
||||
|
||||
LEG_JOINT_PATTERNS = (
|
||||
".*_hip_abduction_joint",
|
||||
".*_hip_pitch_joint",
|
||||
".*_knee_joint",
|
||||
)
|
||||
WHEEL_JOINT_PATTERNS = (".*_wheel_joint",)
|
||||
ALL_JOINT_PATTERNS = LEG_JOINT_PATTERNS + WHEEL_JOINT_PATTERNS
|
||||
|
||||
|
||||
def get_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec.from_file(str(ROBOT_XML))
|
||||
actuators_to_delete = list(spec.actuators)
|
||||
for act in actuators_to_delete:
|
||||
spec.delete(act)
|
||||
return spec
|
||||
|
||||
|
||||
# All 16 actuators have identical physical specs: effort limit = 17 Nm, max velocity = 13 rad/s.
|
||||
# Leg joints: Position PD control (referenced from HIMLoco Go2W: kp = 40, kd = 1).
|
||||
STIFFNESS_LEG = 40.0
|
||||
DAMPING_LEG = 1.0
|
||||
EFFORT_LEG = 17.0
|
||||
|
||||
# Wheel joints: Velocity control with damping = 0.5 and effort limit = 17 Nm.
|
||||
DAMPING_WHEEL = 0.5
|
||||
EFFORT_WHEEL = 17.0
|
||||
|
||||
# Maximum joint speed for all actuators (rad/s)
|
||||
MAX_JOINT_VEL = 13.0
|
||||
|
||||
LEG_ACTUATOR_CFG = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=LEG_JOINT_PATTERNS,
|
||||
stiffness=STIFFNESS_LEG,
|
||||
damping=DAMPING_LEG,
|
||||
effort_limit=EFFORT_LEG,
|
||||
)
|
||||
|
||||
WHEEL_ACTUATOR_CFG = BuiltinVelocityActuatorCfg(
|
||||
target_names_expr=WHEEL_JOINT_PATTERNS,
|
||||
damping=DAMPING_WHEEL,
|
||||
effort_limit=EFFORT_WHEEL,
|
||||
)
|
||||
|
||||
INIT_STATE = EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.40),
|
||||
joint_pos={
|
||||
".*_hip_abduction_joint": 0.0,
|
||||
".*_hip_pitch_joint": 0.9,
|
||||
".*_knee_joint": -1.8,
|
||||
".*_wheel_joint": 0.0,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
|
||||
COLLISION_CFG = CollisionCfg(
|
||||
geom_names_expr=(".*",),
|
||||
contype=0,
|
||||
conaffinity=1,
|
||||
condim={".*_wheel_Link.*": 6, ".*": 1},
|
||||
priority={".*_wheel_Link.*": 1},
|
||||
friction={".*_wheel_Link.*": (0.8, 0.05, 0.01)},
|
||||
)
|
||||
|
||||
ARTICULATION_CFG = EntityArticulationInfoCfg(
|
||||
actuators=(LEG_ACTUATOR_CFG, WHEEL_ACTUATOR_CFG),
|
||||
soft_joint_pos_limit_factor=0.95,
|
||||
)
|
||||
|
||||
|
||||
def get_robot_cfg() -> EntityCfg:
|
||||
return EntityCfg(
|
||||
init_state=INIT_STATE,
|
||||
collisions=(COLLISION_CFG,),
|
||||
spec_fn=get_spec,
|
||||
articulation=ARTICULATION_CFG,
|
||||
)
|
||||
|
||||
|
||||
def get_robot_crawl_cfg() -> EntityCfg:
|
||||
crawl_init_state = EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.20),
|
||||
joint_pos={
|
||||
"(fl|rl)_hip_abduction_joint": 0.4,
|
||||
"(fr|rr)_hip_abduction_joint": -0.4,
|
||||
".*_hip_pitch_joint": 1.65,
|
||||
".*_knee_joint": -2.55,
|
||||
".*_wheel_joint": 0.0,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
return EntityCfg(
|
||||
init_state=crawl_init_state,
|
||||
collisions=(COLLISION_CFG,),
|
||||
spec_fn=get_spec,
|
||||
articulation=ARTICULATION_CFG,
|
||||
)
|
||||
|
||||
|
||||
# Action scale: Leg target position is ±0.25 rad, wheel target velocity is ±10.0 rad/s.
|
||||
# (Policy output range ±1 maps to wheel speeds ±10 rad/s, staying well within the max velocity of 13 rad/s).
|
||||
LEG_POS_SCALE = 0.25
|
||||
WHEEL_VEL_SCALE = 10.0
|
||||
@@ -0,0 +1,5 @@
|
||||
from .modules.him_estimator import HIMEstimator
|
||||
from .modules.him_actor_critic import HIMActorCritic
|
||||
from .algorithms.him_ppo import HIMPPO
|
||||
from .storage.him_rollout_storage import HIMRolloutStorage
|
||||
from .env.vec_env import VecEnv
|
||||
@@ -0,0 +1 @@
|
||||
from .him_ppo import HIMPPO
|
||||
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
|
||||
from ..modules import HIMActorCritic
|
||||
from ..storage import HIMRolloutStorage
|
||||
|
||||
|
||||
class HIMPPO:
|
||||
actor_critic: HIMActorCritic
|
||||
|
||||
def __init__(self,
|
||||
actor_critic,
|
||||
num_learning_epochs=1,
|
||||
num_mini_batches=1,
|
||||
clip_param=0.2,
|
||||
gamma=0.998,
|
||||
lam=0.95,
|
||||
value_loss_coef=1.0,
|
||||
entropy_coef=0.0,
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
device='cpu',
|
||||
):
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
|
||||
# PPO components
|
||||
self.actor_critic = actor_critic
|
||||
self.actor_critic.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
self.optimizer = torch.optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
|
||||
self.transition = HIMRolloutStorage.Transition()
|
||||
|
||||
# PPO parameters
|
||||
self.clip_param = clip_param
|
||||
self.num_learning_epochs = num_learning_epochs
|
||||
self.num_mini_batches = num_mini_batches
|
||||
self.value_loss_coef = value_loss_coef
|
||||
self.entropy_coef = entropy_coef
|
||||
self.gamma = gamma
|
||||
self.lam = lam
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.use_clipped_value_loss = use_clipped_value_loss
|
||||
|
||||
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
|
||||
self.storage = HIMRolloutStorage(
|
||||
num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device
|
||||
)
|
||||
|
||||
def test_mode(self):
|
||||
self.actor_critic.test()
|
||||
|
||||
def train_mode(self):
|
||||
self.actor_critic.train()
|
||||
|
||||
def act(self, obs, critic_obs):
|
||||
self.transition.actions = self.actor_critic.act(obs).detach()
|
||||
self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
|
||||
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(
|
||||
self.transition.actions
|
||||
).detach()
|
||||
self.transition.action_mean = self.actor_critic.action_mean.detach()
|
||||
self.transition.action_sigma = self.actor_critic.action_std.detach()
|
||||
self.transition.observations = obs
|
||||
self.transition.critic_observations = critic_obs
|
||||
return self.transition.actions
|
||||
|
||||
def process_env_step(self, rewards, dones, infos, next_critic_obs):
|
||||
self.transition.next_critic_observations = next_critic_obs.clone()
|
||||
self.transition.rewards = rewards.clone()
|
||||
self.transition.dones = dones
|
||||
# Bootstrapping on time outs
|
||||
if 'time_outs' in infos:
|
||||
self.transition.rewards += self.gamma * torch.squeeze(
|
||||
self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1
|
||||
)
|
||||
|
||||
self.storage.add_transitions(self.transition)
|
||||
self.transition.clear()
|
||||
self.actor_critic.reset(dones)
|
||||
|
||||
def compute_returns(self, last_critic_obs):
|
||||
last_values = self.actor_critic.evaluate(last_critic_obs).detach()
|
||||
self.storage.compute_returns(last_values, self.gamma, self.lam)
|
||||
|
||||
def update(self):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_estimation_loss = 0
|
||||
mean_swap_loss = 0
|
||||
|
||||
generator = self.storage.mini_batch_generator(
|
||||
self.num_mini_batches, self.num_learning_epochs
|
||||
)
|
||||
|
||||
for obs_batch, critic_obs_batch, actions_batch, next_critic_obs_batch, \
|
||||
target_values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch in generator:
|
||||
|
||||
self.actor_critic.act(obs_batch)
|
||||
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.actor_critic.evaluate(critic_obs_batch)
|
||||
mu_batch = self.actor_critic.action_mean
|
||||
sigma_batch = self.actor_critic.action_std
|
||||
entropy_batch = self.actor_critic.entropy
|
||||
|
||||
# KL
|
||||
if self.desired_kl is not None and self.schedule == 'adaptive':
|
||||
with torch.inference_mode():
|
||||
kl = torch.sum(
|
||||
torch.log(sigma_batch / old_sigma_batch + 1.e-5)
|
||||
+ (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch))
|
||||
/ (2.0 * torch.square(sigma_batch)) - 0.5,
|
||||
axis=-1,
|
||||
)
|
||||
kl_mean = torch.mean(kl)
|
||||
|
||||
if kl_mean > self.desired_kl * 2.0:
|
||||
self.learning_rate = max(1e-5, self.learning_rate / 1.5)
|
||||
elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
|
||||
self.learning_rate = min(1e-2, self.learning_rate * 1.5)
|
||||
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
# Estimator Update
|
||||
estimation_loss, swap_loss = self.actor_critic.estimator.update(
|
||||
obs_batch, next_critic_obs_batch, lr=self.learning_rate
|
||||
)
|
||||
|
||||
# Surrogate loss
|
||||
ratio = torch.exp(
|
||||
actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)
|
||||
)
|
||||
surrogate = -torch.squeeze(advantages_batch) * ratio
|
||||
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
|
||||
ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
|
||||
)
|
||||
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
|
||||
|
||||
# Value function loss
|
||||
if self.use_clipped_value_loss:
|
||||
value_clipped = target_values_batch + (
|
||||
value_batch - target_values_batch
|
||||
).clamp(-self.clip_param, self.clip_param)
|
||||
value_losses = (value_batch - returns_batch).pow(2)
|
||||
value_losses_clipped = (value_clipped - returns_batch).pow(2)
|
||||
value_loss = torch.max(value_losses, value_losses_clipped).mean()
|
||||
else:
|
||||
value_loss = (returns_batch - value_batch).pow(2).mean()
|
||||
|
||||
loss = (
|
||||
surrogate_loss
|
||||
+ self.value_loss_coef * value_loss
|
||||
- self.entropy_coef * entropy_batch.mean()
|
||||
)
|
||||
|
||||
# Gradient step
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_estimation_loss += estimation_loss
|
||||
mean_swap_loss += swap_loss
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_estimation_loss /= num_updates
|
||||
mean_swap_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, mean_estimation_loss, mean_swap_loss
|
||||
|
||||
|
||||
# Import nn here to avoid circular dependency at module level
|
||||
import torch.nn as nn
|
||||
@@ -0,0 +1 @@
|
||||
from .vec_env import VecEnv
|
||||
@@ -0,0 +1,40 @@
|
||||
import torch
|
||||
|
||||
|
||||
class VecEnv:
|
||||
"""Abstract vectorized environment interface for HIMLoco RSL-RL."""
|
||||
|
||||
num_envs: int
|
||||
num_obs: int | None = None
|
||||
num_one_step_obs: int | None = None
|
||||
num_privileged_obs: int | None = None
|
||||
num_one_step_privileged_obs: int | None = None
|
||||
num_actions: int
|
||||
max_episode_length: int
|
||||
device: str
|
||||
|
||||
def get_observations(self) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_privileged_observations(self) -> torch.Tensor | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def step(self, actions: torch.Tensor) -> tuple:
|
||||
raise NotImplementedError
|
||||
|
||||
def reset(self) -> tuple:
|
||||
raise NotImplementedError
|
||||
|
||||
def seed(self, seed: int = -1) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
def close(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def episode_length_buf(self) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
|
||||
@episode_length_buf.setter
|
||||
def episode_length_buf(self, value: torch.Tensor) -> None:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,2 @@
|
||||
from .him_estimator import HIMEstimator
|
||||
from .him_actor_critic import HIMActorCritic
|
||||
@@ -0,0 +1,206 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
from ..modules.him_estimator import HIMEstimator
|
||||
|
||||
|
||||
class RunningMeanStd:
|
||||
def __init__(self, shape, device):
|
||||
self.n = 1e-4
|
||||
self.uninitialized = True
|
||||
self.mean = torch.zeros(shape, device=device)
|
||||
self.var = torch.ones(shape, device=device)
|
||||
|
||||
def update(self, x):
|
||||
count = self.n
|
||||
batch_count = x.size(0)
|
||||
tot_count = count + batch_count
|
||||
|
||||
old_mean = self.mean.clone()
|
||||
delta = torch.mean(x, dim=0) - old_mean
|
||||
|
||||
self.mean = old_mean + delta * batch_count / tot_count
|
||||
m_a = self.var * count
|
||||
m_b = x.var(dim=0) * batch_count
|
||||
M2 = m_a + m_b + torch.square(delta) * count * batch_count / tot_count
|
||||
self.var = M2 / tot_count
|
||||
self.n = tot_count
|
||||
|
||||
|
||||
class Normalization:
|
||||
def __init__(self, shape, device='cuda:0'):
|
||||
self.running_ms = RunningMeanStd(shape=shape, device=device)
|
||||
|
||||
def __call__(self, x, update=False):
|
||||
if update:
|
||||
self.running_ms.update(x)
|
||||
x = (x - self.running_ms.mean) / (torch.sqrt(self.running_ms.var) + 1e-4)
|
||||
return x
|
||||
|
||||
|
||||
class HIMActorCritic(nn.Module):
|
||||
is_recurrent = False
|
||||
|
||||
def __init__(self, num_actor_obs,
|
||||
num_critic_obs,
|
||||
num_one_step_obs,
|
||||
num_actions,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation='elu',
|
||||
init_noise_std=1.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCritic.__init__ got unexpected arguments, which will be ignored: "
|
||||
+ str([key for key in kwargs.keys()]))
|
||||
super(HIMActorCritic, self).__init__()
|
||||
|
||||
activation = get_activation(activation)
|
||||
|
||||
self.history_size = int(num_actor_obs / num_one_step_obs)
|
||||
self.num_actor_obs = num_actor_obs
|
||||
self.num_actions = num_actions
|
||||
self.num_one_step_obs = num_one_step_obs
|
||||
|
||||
mlp_input_dim_a = num_one_step_obs + 3 + 16
|
||||
mlp_input_dim_c = num_critic_obs
|
||||
|
||||
# Estimator
|
||||
self.estimator = HIMEstimator(
|
||||
temporal_steps=self.history_size,
|
||||
num_one_step_obs=num_one_step_obs,
|
||||
)
|
||||
|
||||
# Policy
|
||||
actor_layers = []
|
||||
actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
|
||||
actor_layers.append(activation)
|
||||
for l in range(len(actor_hidden_dims)):
|
||||
if l == len(actor_hidden_dims) - 1:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
|
||||
else:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
|
||||
actor_layers.append(activation)
|
||||
self.actor = nn.Sequential(*actor_layers)
|
||||
|
||||
# Value function
|
||||
critic_layers = []
|
||||
critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
|
||||
critic_layers.append(activation)
|
||||
for l in range(len(critic_hidden_dims)):
|
||||
if l == len(critic_hidden_dims) - 1:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
|
||||
else:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
|
||||
critic_layers.append(activation)
|
||||
self.critic = nn.Sequential(*critic_layers)
|
||||
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
print(f'Estimator: {self.estimator.encoder}')
|
||||
|
||||
# Action noise
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
@staticmethod
|
||||
def init_weights(sequential, scales):
|
||||
[torch.nn.init.orthogonal_(module.weight, gain=scales[idx])
|
||||
for idx, module in
|
||||
enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]
|
||||
|
||||
def reset(self, dones=None):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, obs_history):
|
||||
with torch.no_grad():
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actor_input = torch.cat(
|
||||
(obs_history[:, :self.num_one_step_obs], vel, latent), dim=-1
|
||||
)
|
||||
|
||||
if torch.isnan(actor_input).any() or torch.isinf(actor_input).any():
|
||||
print(f"[ERROR] NaN/Inf detected in actor_input before normalization!")
|
||||
print(f" - obs_history: min={obs_history.min():.4f}, max={obs_history.max():.4f}")
|
||||
print(f" - vel: min={vel.min():.4f}, max={vel.max():.4f}")
|
||||
print(f" - latent: min={latent.min():.4f}, max={latent.max():.4f}")
|
||||
raise ValueError("NaN/Inf in actor_input before normalization")
|
||||
|
||||
mean = self.actor(actor_input)
|
||||
|
||||
if torch.isnan(mean).any() or torch.isinf(mean).any():
|
||||
print(f"[ERROR] NaN/Inf in actor output (mean)!")
|
||||
for name, param in self.actor.named_parameters():
|
||||
if torch.isnan(param).any():
|
||||
print(f" - NaN in actor weight: {name}")
|
||||
if torch.isinf(param).any():
|
||||
print(f" - Inf in actor weight: {name}")
|
||||
raise ValueError("NaN or Inf detected in actor network output!")
|
||||
|
||||
self.distribution = Normal(mean, mean * 0. + self.std)
|
||||
|
||||
def act(self, obs_history=None, **kwargs):
|
||||
self.update_distribution(obs_history)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, obs_history, observations=None):
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actions_mean = self.actor(
|
||||
torch.cat(
|
||||
(obs_history[:, :self.num_one_step_obs], vel, latent), dim=-1
|
||||
)
|
||||
)
|
||||
return actions_mean
|
||||
|
||||
def test_inference(self, obs_history, observations=None):
|
||||
vel, latent = self.estimator(obs_history)
|
||||
actions_mean = self.actor(
|
||||
torch.cat(
|
||||
(obs_history[:, :self.num_one_step_obs], vel, latent), dim=-1
|
||||
)
|
||||
)
|
||||
estimator_output = torch.cat((vel, latent), dim=-1)
|
||||
return actions_mean, estimator_output
|
||||
|
||||
def evaluate(self, critic_observations, **kwargs):
|
||||
value = self.critic(critic_observations)
|
||||
return value
|
||||
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
@@ -0,0 +1,153 @@
|
||||
import copy
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class HIMEstimator(nn.Module):
|
||||
def __init__(self,
|
||||
temporal_steps,
|
||||
num_one_step_obs,
|
||||
enc_hidden_dims=[128, 64, 16],
|
||||
tar_hidden_dims=[128, 64],
|
||||
activation='elu',
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=10.0,
|
||||
num_prototype=32,
|
||||
temperature=3.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("Estimator.__init__ got unexpected arguments, which will be ignored: "
|
||||
+ str([key for key in kwargs.keys()]))
|
||||
super(HIMEstimator, self).__init__()
|
||||
activation = get_activation(activation)
|
||||
|
||||
self.temporal_steps = temporal_steps
|
||||
self.num_one_step_obs = num_one_step_obs
|
||||
self.num_latent = enc_hidden_dims[-1]
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.temperature = temperature
|
||||
|
||||
# Encoder
|
||||
enc_input_dim = self.temporal_steps * self.num_one_step_obs
|
||||
enc_layers = []
|
||||
for l in range(len(enc_hidden_dims) - 1):
|
||||
enc_layers += [nn.Linear(enc_input_dim, enc_hidden_dims[l]), activation]
|
||||
enc_input_dim = enc_hidden_dims[l]
|
||||
enc_layers += [nn.Linear(enc_input_dim, enc_hidden_dims[-1] + 3)]
|
||||
self.encoder = nn.Sequential(*enc_layers)
|
||||
|
||||
# Target
|
||||
tar_input_dim = self.num_one_step_obs
|
||||
tar_layers = []
|
||||
for l in range(len(tar_hidden_dims)):
|
||||
tar_layers += [nn.Linear(tar_input_dim, tar_hidden_dims[l]), activation]
|
||||
tar_input_dim = tar_hidden_dims[l]
|
||||
tar_layers += [nn.Linear(tar_input_dim, enc_hidden_dims[-1])]
|
||||
self.target = nn.Sequential(*tar_layers)
|
||||
|
||||
# Prototype
|
||||
self.proto = nn.Embedding(num_prototype, enc_hidden_dims[-1])
|
||||
|
||||
# Optimizer
|
||||
self.learning_rate = learning_rate
|
||||
self.optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
|
||||
|
||||
def get_latent(self, obs_history):
|
||||
vel, z = self.encode(obs_history)
|
||||
return vel.detach(), z.detach()
|
||||
|
||||
def forward(self, obs_history):
|
||||
parts = self.encoder(obs_history.detach())
|
||||
vel, z = parts[..., :3], parts[..., 3:]
|
||||
z = F.normalize(z, dim=-1, p=2)
|
||||
return vel.detach(), z.detach()
|
||||
|
||||
def encode(self, obs_history):
|
||||
parts = self.encoder(obs_history.detach())
|
||||
vel, z = parts[..., :3], parts[..., 3:]
|
||||
z = F.normalize(z, dim=-1, p=2)
|
||||
return vel, z
|
||||
|
||||
def update(self, obs_history, next_critic_obs, lr=None):
|
||||
if lr is not None:
|
||||
self.learning_rate = lr
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
vel = next_critic_obs[:, self.num_one_step_obs:self.num_one_step_obs + 3].detach()
|
||||
next_obs = next_critic_obs.detach()[:, 3:self.num_one_step_obs + 3]
|
||||
|
||||
z_s = self.encoder(obs_history)
|
||||
z_t = self.target(next_obs)
|
||||
pred_vel, z_s = z_s[..., :3], z_s[..., 3:]
|
||||
|
||||
z_s = F.normalize(z_s, dim=-1, p=2)
|
||||
z_t = F.normalize(z_t, dim=-1, p=2)
|
||||
|
||||
with torch.no_grad():
|
||||
w = self.proto.weight.data.clone()
|
||||
w = F.normalize(w, dim=-1, p=2)
|
||||
self.proto.weight.copy_(w)
|
||||
|
||||
score_s = z_s @ self.proto.weight.T
|
||||
score_t = z_t @ self.proto.weight.T
|
||||
|
||||
with torch.no_grad():
|
||||
q_s = sinkhorn(score_s)
|
||||
q_t = sinkhorn(score_t)
|
||||
|
||||
log_p_s = F.log_softmax(score_s / self.temperature, dim=-1)
|
||||
log_p_t = F.log_softmax(score_t / self.temperature, dim=-1)
|
||||
|
||||
swap_loss = -0.5 * (q_s * log_p_t + q_t * log_p_s).mean()
|
||||
estimation_loss = F.mse_loss(pred_vel, vel)
|
||||
losses = estimation_loss + swap_loss
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
losses.backward()
|
||||
nn.utils.clip_grad_norm_(self.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
return estimation_loss.item(), swap_loss.item()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sinkhorn(out, eps=0.05, iters=3):
|
||||
Q = torch.exp(out / eps).T
|
||||
K, B = Q.shape[0], Q.shape[1]
|
||||
Q /= Q.sum()
|
||||
|
||||
for it in range(iters):
|
||||
# normalize each row: total weight per prototype must be 1/K
|
||||
Q /= torch.sum(Q, dim=1, keepdim=True)
|
||||
Q /= K
|
||||
|
||||
# normalize each column: total weight per sample must be 1/B
|
||||
Q /= torch.sum(Q, dim=0, keepdim=True)
|
||||
Q /= B
|
||||
return (Q * B).T
|
||||
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "silu":
|
||||
return nn.SiLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
@@ -0,0 +1,305 @@
|
||||
import time
|
||||
import os
|
||||
import statistics
|
||||
from collections import deque
|
||||
|
||||
import torch
|
||||
|
||||
from ..algorithms import HIMPPO
|
||||
from ..modules import HIMActorCritic
|
||||
from ..env import VecEnv
|
||||
|
||||
|
||||
class HIMOnPolicyRunner:
|
||||
"""On-policy runner for HIMLoco training.
|
||||
|
||||
Uses the HIM architecture: asymmetric actor-critic with history-based
|
||||
estimator for privileged information extraction.
|
||||
"""
|
||||
|
||||
def __init__(self, env: VecEnv, train_cfg, log_dir=None, device='cpu'):
|
||||
self.cfg = train_cfg
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
|
||||
# Determine observation dimensions
|
||||
if self.env.num_privileged_obs is not None:
|
||||
num_critic_obs = self.env.num_privileged_obs
|
||||
else:
|
||||
num_critic_obs = self.env.num_obs
|
||||
self.num_actor_obs = self.env.num_obs
|
||||
self.num_critic_obs = num_critic_obs
|
||||
|
||||
# Initialize policy network
|
||||
actor_critic_class = eval(train_cfg["policy_class_name"])
|
||||
actor_critic: HIMActorCritic = actor_critic_class(
|
||||
self.env.num_obs,
|
||||
num_critic_obs,
|
||||
self.env.num_one_step_obs,
|
||||
self.env.num_actions,
|
||||
**self.policy_cfg
|
||||
).to(self.device)
|
||||
|
||||
# Initialize algorithm
|
||||
alg_class = eval(train_cfg["algorithm_class_name"])
|
||||
self.alg: HIMPPO = alg_class(actor_critic, device=self.device, **self.alg_cfg)
|
||||
|
||||
# Training configuration
|
||||
self.num_steps_per_env = train_cfg["num_steps_per_env"]
|
||||
self.save_interval = train_cfg["save_interval"]
|
||||
|
||||
# Initialize storage
|
||||
self.alg.init_storage(
|
||||
self.env.num_envs,
|
||||
self.num_steps_per_env,
|
||||
[self.env.num_obs],
|
||||
[self.env.num_privileged_obs],
|
||||
[self.env.num_actions]
|
||||
)
|
||||
|
||||
# Logging
|
||||
self.log_dir = log_dir
|
||||
self.writer = None
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.current_learning_iteration = 0
|
||||
self.logger_type = self.cfg.get("logger", "tensorboard").lower()
|
||||
|
||||
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
|
||||
if self.log_dir is not None and self.writer is None:
|
||||
if self.logger_type == "neptune":
|
||||
from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter
|
||||
self.writer = NeptuneSummaryWriter(
|
||||
log_dir=self.log_dir, flush_secs=10, cfg=self.alg_cfg
|
||||
)
|
||||
elif self.logger_type == "wandb":
|
||||
from rsl_rl.utils.wandb_utils import WandbSummaryWriter
|
||||
self.writer = WandbSummaryWriter(
|
||||
log_dir=self.log_dir, flush_secs=10, cfg=self.cfg
|
||||
)
|
||||
if hasattr(self.writer, "store_config"):
|
||||
env_cfg = getattr(self.env, "cfg", None)
|
||||
if env_cfg is None and hasattr(self.env, "unwrapped"):
|
||||
env_cfg = getattr(self.env.unwrapped, "cfg", None)
|
||||
self.writer.store_config(env_cfg, self.cfg)
|
||||
elif self.logger_type == "tensorboard":
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Logger type not found. Choose 'neptune', 'wandb' or 'tensorboard'."
|
||||
)
|
||||
|
||||
if init_at_random_ep_len:
|
||||
self.env.episode_length_buf = torch.randint_like(
|
||||
self.env.episode_length_buf,
|
||||
high=int(self.env.max_episode_length)
|
||||
)
|
||||
obs = self.env.get_observations()
|
||||
privileged_obs = self.env.get_privileged_observations()
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs = obs.to(self.device), critic_obs.to(self.device)
|
||||
self.alg.actor_critic.train()
|
||||
|
||||
ep_infos = []
|
||||
rewbuffer = deque(maxlen=100)
|
||||
lenbuffer = deque(maxlen=100)
|
||||
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
start_iter = self.current_learning_iteration
|
||||
tot_iter = start_iter + num_learning_iterations
|
||||
for it in range(start_iter, tot_iter):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for i in range(self.num_steps_per_env):
|
||||
actions = self.alg.act(obs, critic_obs)
|
||||
obs, privileged_obs, rewards, dones, infos, \
|
||||
termination_ids, termination_privileged_obs = self.env.step(actions)
|
||||
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs = obs.to(self.device)
|
||||
critic_obs = critic_obs.to(self.device)
|
||||
rewards = rewards.to(self.device)
|
||||
dones = dones.to(self.device)
|
||||
termination_ids = termination_ids.to(self.device)
|
||||
termination_privileged_obs = termination_privileged_obs.to(self.device)
|
||||
|
||||
next_critic_obs = critic_obs.clone().detach()
|
||||
next_critic_obs[termination_ids] = termination_privileged_obs.clone().detach()
|
||||
|
||||
self.alg.process_env_step(rewards, dones, infos, next_critic_obs)
|
||||
|
||||
if self.log_dir is not None:
|
||||
if 'episode' in infos:
|
||||
ep_infos.append(infos['episode'])
|
||||
elif 'log' in infos:
|
||||
ep_infos.append(infos['log'])
|
||||
cur_reward_sum += rewards
|
||||
cur_episode_length += 1
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
rewbuffer.extend(
|
||||
cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()
|
||||
)
|
||||
lenbuffer.extend(
|
||||
cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()
|
||||
)
|
||||
cur_reward_sum[new_ids] = 0
|
||||
cur_episode_length[new_ids] = 0
|
||||
|
||||
stop = time.time()
|
||||
collection_time = stop - start
|
||||
|
||||
# Learning step
|
||||
start = stop
|
||||
self.alg.compute_returns(critic_obs)
|
||||
|
||||
mean_value_loss, mean_surrogate_loss, mean_estimation_loss, mean_swap_loss = \
|
||||
self.alg.update()
|
||||
stop = time.time()
|
||||
learn_time = stop - start
|
||||
self.current_learning_iteration = it
|
||||
|
||||
if self.log_dir is not None:
|
||||
self.log(locals())
|
||||
if it % self.save_interval == 0:
|
||||
self.save(os.path.join(self.log_dir, f'model_{it}.pt'))
|
||||
ep_infos.clear()
|
||||
|
||||
if self.log_dir is not None:
|
||||
self.save(os.path.join(
|
||||
self.log_dir, f'model_{self.current_learning_iteration}.pt'
|
||||
))
|
||||
|
||||
def log(self, locs, width=80, pad=35):
|
||||
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
|
||||
self.tot_time += locs['collection_time'] + locs['learn_time']
|
||||
iteration_time = locs['collection_time'] + locs['learn_time']
|
||||
|
||||
ep_string = ""
|
||||
if locs['ep_infos']:
|
||||
for key in locs['ep_infos'][0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
for ep_info in locs['ep_infos']:
|
||||
if key not in ep_info:
|
||||
continue
|
||||
if not isinstance(ep_info[key], torch.Tensor):
|
||||
ep_info[key] = torch.Tensor([ep_info[key]])
|
||||
if len(ep_info[key].shape) == 0:
|
||||
ep_info[key] = ep_info[key].unsqueeze(0)
|
||||
infotensor = torch.cat(
|
||||
(infotensor, ep_info[key].to(self.device))
|
||||
)
|
||||
if infotensor.numel() > 0:
|
||||
value = torch.mean(infotensor)
|
||||
if "/" in key:
|
||||
self.writer.add_scalar(key, value, locs['it'])
|
||||
ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n"""
|
||||
else:
|
||||
self.writer.add_scalar('Episode/' + key, value, locs['it'])
|
||||
ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
|
||||
|
||||
mean_std = self.alg.actor_critic.std.mean()
|
||||
fps = int(
|
||||
self.num_steps_per_env * self.env.num_envs
|
||||
/ (locs['collection_time'] + locs['learn_time'])
|
||||
)
|
||||
|
||||
loss_dict = {
|
||||
'value_function': locs['mean_value_loss'],
|
||||
'surrogate': locs['mean_surrogate_loss'],
|
||||
'estimation': locs['mean_estimation_loss'],
|
||||
'swap': locs['mean_swap_loss'],
|
||||
}
|
||||
for key, value in loss_dict.items():
|
||||
self.writer.add_scalar(f'Loss/{key}', value, locs['it'])
|
||||
|
||||
self.writer.add_scalar('Loss/learning_rate', self.alg.learning_rate, locs['it'])
|
||||
self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), locs['it'])
|
||||
self.writer.add_scalar('Perf/total_fps', fps, locs['it'])
|
||||
self.writer.add_scalar('Perf/collection_time', locs['collection_time'], locs['it'])
|
||||
self.writer.add_scalar('Perf/learning_time', locs['learn_time'], locs['it'])
|
||||
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
self.writer.add_scalar(
|
||||
'Train/mean_reward', statistics.mean(locs['rewbuffer']), locs['it']
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
'Train/mean_episode_length', statistics.mean(locs['lenbuffer']), locs['it']
|
||||
)
|
||||
if self.logger_type != "wandb":
|
||||
self.writer.add_scalar(
|
||||
'Train/mean_reward/time', statistics.mean(locs['rewbuffer']),
|
||||
self.tot_time
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
'Train/mean_episode_length/time',
|
||||
statistics.mean(locs['lenbuffer']), self.tot_time
|
||||
)
|
||||
|
||||
str_header = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m "
|
||||
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
log_string = (
|
||||
f"""{'#' * width}\n"""
|
||||
f"""{str_header.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs['collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
)
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n"""
|
||||
log_string += (
|
||||
f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n"""
|
||||
f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n"""
|
||||
)
|
||||
else:
|
||||
log_string = (
|
||||
f"""{'#' * width}\n"""
|
||||
f"""{str_header.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs['collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
)
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n"""
|
||||
|
||||
log_string += ep_string
|
||||
log_string += (
|
||||
f"""{'-' * width}\n"""
|
||||
f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n"""
|
||||
f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n"""
|
||||
f"""{'Time elapsed:':>{pad}} {time.strftime('%H:%M:%S', time.gmtime(self.tot_time))}\n"""
|
||||
)
|
||||
print(log_string)
|
||||
|
||||
def save(self, path, infos=None):
|
||||
saved_dict = {
|
||||
'model_state_dict': self.alg.actor_critic.state_dict(),
|
||||
'optimizer_state_dict': self.alg.optimizer.state_dict(),
|
||||
'estimator_optimizer_state_dict':
|
||||
self.alg.actor_critic.estimator.optimizer.state_dict(),
|
||||
'iter': self.current_learning_iteration,
|
||||
'infos': infos,
|
||||
}
|
||||
torch.save(saved_dict, path)
|
||||
if self.cfg.get("upload_model", False) and hasattr(self.writer, "save_model"):
|
||||
self.writer.save_model(path, self.current_learning_iteration)
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
loaded_dict = torch.load(path, map_location=self.device, weights_only=False)
|
||||
self.alg.actor_critic.load_state_dict(loaded_dict['model_state_dict'])
|
||||
if load_optimizer:
|
||||
self.alg.optimizer.load_state_dict(loaded_dict['optimizer_state_dict'])
|
||||
self.alg.actor_critic.estimator.optimizer.load_state_dict(
|
||||
loaded_dict['estimator_optimizer_state_dict']
|
||||
)
|
||||
self.current_learning_iteration = loaded_dict['iter']
|
||||
return loaded_dict['infos']
|
||||
|
||||
def get_inference_policy(self, device=None):
|
||||
self.alg.actor_critic.eval()
|
||||
if device is not None:
|
||||
self.alg.actor_critic.to(device)
|
||||
return self.alg.actor_critic.act_inference
|
||||
@@ -0,0 +1 @@
|
||||
from .him_rollout_storage import HIMRolloutStorage
|
||||
@@ -0,0 +1,170 @@
|
||||
import torch
|
||||
|
||||
|
||||
class HIMRolloutStorage:
|
||||
class Transition:
|
||||
def __init__(self):
|
||||
self.observations = None
|
||||
self.critic_observations = None
|
||||
self.actions = None
|
||||
self.rewards = None
|
||||
self.dones = None
|
||||
self.values = None
|
||||
self.actions_log_prob = None
|
||||
self.action_mean = None
|
||||
self.action_sigma = None
|
||||
self.next_critic_observations = None
|
||||
|
||||
def clear(self):
|
||||
self.__init__()
|
||||
|
||||
def __init__(self, num_envs, num_transitions_per_env, obs_shape,
|
||||
privileged_obs_shape, actions_shape, device='cpu'):
|
||||
self.device = device
|
||||
|
||||
self.obs_shape = obs_shape
|
||||
self.privileged_obs_shape = privileged_obs_shape
|
||||
self.actions_shape = actions_shape
|
||||
|
||||
# Core
|
||||
self.observations = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *obs_shape, device=self.device
|
||||
)
|
||||
if privileged_obs_shape[0] is not None:
|
||||
self.privileged_observations = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device
|
||||
)
|
||||
self.next_privileged_observations = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device
|
||||
)
|
||||
else:
|
||||
self.privileged_observations = None
|
||||
self.next_privileged_observations = None
|
||||
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.actions = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *actions_shape, device=self.device
|
||||
)
|
||||
self.dones = torch.zeros(
|
||||
num_transitions_per_env, num_envs, 1, device=self.device
|
||||
).byte()
|
||||
|
||||
# For PPO
|
||||
self.actions_log_prob = torch.zeros(
|
||||
num_transitions_per_env, num_envs, 1, device=self.device
|
||||
)
|
||||
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.advantages = torch.zeros(
|
||||
num_transitions_per_env, num_envs, 1, device=self.device
|
||||
)
|
||||
self.mu = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *actions_shape, device=self.device
|
||||
)
|
||||
self.sigma = torch.zeros(
|
||||
num_transitions_per_env, num_envs, *actions_shape, device=self.device
|
||||
)
|
||||
|
||||
self.num_transitions_per_env = num_transitions_per_env
|
||||
self.num_envs = num_envs
|
||||
|
||||
self.step = 0
|
||||
|
||||
def add_transitions(self, transition: Transition):
|
||||
if self.step >= self.num_transitions_per_env:
|
||||
raise AssertionError("Rollout buffer overflow")
|
||||
self.observations[self.step].copy_(transition.observations)
|
||||
if self.privileged_observations is not None:
|
||||
self.privileged_observations[self.step].copy_(transition.critic_observations)
|
||||
if self.next_privileged_observations is not None:
|
||||
self.next_privileged_observations[self.step].copy_(
|
||||
transition.next_critic_observations
|
||||
)
|
||||
self.actions[self.step].copy_(transition.actions)
|
||||
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
|
||||
self.dones[self.step].copy_(transition.dones.view(-1, 1))
|
||||
self.values[self.step].copy_(transition.values)
|
||||
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
|
||||
self.mu[self.step].copy_(transition.action_mean)
|
||||
self.sigma[self.step].copy_(transition.action_sigma)
|
||||
self.step += 1
|
||||
|
||||
def clear(self):
|
||||
self.step = 0
|
||||
|
||||
def compute_returns(self, last_values, gamma, lam):
|
||||
advantage = 0
|
||||
for step in reversed(range(self.num_transitions_per_env)):
|
||||
if step == self.num_transitions_per_env - 1:
|
||||
next_values = last_values
|
||||
else:
|
||||
next_values = self.values[step + 1]
|
||||
next_is_not_terminal = 1.0 - self.dones[step].float()
|
||||
delta = (
|
||||
self.rewards[step]
|
||||
+ next_is_not_terminal * gamma * next_values
|
||||
- self.values[step]
|
||||
)
|
||||
advantage = delta + next_is_not_terminal * gamma * lam * advantage
|
||||
self.returns[step] = advantage + self.values[step]
|
||||
|
||||
self.advantages = self.returns - self.values
|
||||
self.advantages = (self.advantages - self.advantages.mean()) / (
|
||||
self.advantages.std() + 1e-8
|
||||
)
|
||||
|
||||
def get_statistics(self):
|
||||
done = self.dones
|
||||
done[-1] = 1
|
||||
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
|
||||
done_indices = torch.cat(
|
||||
(flat_dones.new_tensor([-1], dtype=torch.int64),
|
||||
flat_dones.nonzero(as_tuple=False)[:, 0])
|
||||
)
|
||||
trajectory_lengths = done_indices[1:] - done_indices[:-1]
|
||||
return trajectory_lengths.float().mean(), self.rewards.mean()
|
||||
|
||||
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
batch_size = self.num_envs * self.num_transitions_per_env
|
||||
mini_batch_size = batch_size // num_mini_batches
|
||||
indices = torch.randperm(
|
||||
num_mini_batches * mini_batch_size, requires_grad=False, device=self.device
|
||||
)
|
||||
|
||||
observations = self.observations.flatten(0, 1)
|
||||
if self.privileged_observations is not None:
|
||||
critic_observations = self.privileged_observations.flatten(0, 1)
|
||||
next_critic_observations = self.next_privileged_observations.flatten(0, 1)
|
||||
else:
|
||||
critic_observations = observations
|
||||
next_critic_observations = observations
|
||||
|
||||
actions = self.actions.flatten(0, 1)
|
||||
values = self.values.flatten(0, 1)
|
||||
returns = self.returns.flatten(0, 1)
|
||||
old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
|
||||
advantages = self.advantages.flatten(0, 1)
|
||||
old_mu = self.mu.flatten(0, 1)
|
||||
old_sigma = self.sigma.flatten(0, 1)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
start = i * mini_batch_size
|
||||
end = (i + 1) * mini_batch_size
|
||||
batch_idx = indices[start:end]
|
||||
|
||||
obs_batch = observations[batch_idx]
|
||||
next_critic_observations_batch = next_critic_observations[batch_idx]
|
||||
critic_observations_batch = critic_observations[batch_idx]
|
||||
actions_batch = actions[batch_idx]
|
||||
target_values_batch = values[batch_idx]
|
||||
returns_batch = returns[batch_idx]
|
||||
old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
|
||||
advantages_batch = advantages[batch_idx]
|
||||
old_mu_batch = old_mu[batch_idx]
|
||||
old_sigma_batch = old_sigma[batch_idx]
|
||||
yield (
|
||||
obs_batch, critic_observations_batch, actions_batch,
|
||||
next_critic_observations_batch, target_values_batch,
|
||||
advantages_batch, returns_batch,
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
|
||||
)
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Wrapper to adapt mjlab HimManagerBasedRlEnv for HIMLoco RSL-RL.
|
||||
|
||||
Handles:
|
||||
- Observation history stacking (actor obs get history, critic gets single step)
|
||||
- Termination privileged observation extraction
|
||||
- 7-value step return: (obs, privileged_obs, rewards, dones, infos,
|
||||
termination_ids, termination_privileged_obs)
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from him_mjlab.envs.him_manager_based_rl_env import HimManagerBasedRlEnv
|
||||
from ..env.vec_env import VecEnv
|
||||
|
||||
|
||||
class HimMjlabVecEnvWrapper(VecEnv):
|
||||
"""Wraps mjlab HimManagerBasedRlEnv for HIMLoco RSL-RL.
|
||||
|
||||
Key features:
|
||||
- Converts dict observations to flat tensors
|
||||
- Maintains history buffers for actor observations
|
||||
- Tracks terminated envs and their pre-reset privileged observations
|
||||
- Returns 7-value step required by HIMOnPolicyRunner
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
env: HimManagerBasedRlEnv,
|
||||
history_length: int = 5,
|
||||
privileged_history_length: int = 0,
|
||||
):
|
||||
if not isinstance(env.unwrapped, HimManagerBasedRlEnv):
|
||||
raise ValueError(
|
||||
"The environment must be HimManagerBasedRlEnv. "
|
||||
f"Got: {type(env)}"
|
||||
)
|
||||
|
||||
self.env = env
|
||||
self.reorder_indices = [0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
|
||||
|
||||
self.num_envs = self.unwrapped.num_envs
|
||||
self.device = self.unwrapped.device
|
||||
self.max_episode_length = self.unwrapped.max_episode_length
|
||||
self.num_actions = self.unwrapped.action_manager.total_action_dim
|
||||
|
||||
# Single-step observation dimensions from observation manager groups
|
||||
self.num_one_step_obs = self.unwrapped.observation_manager.group_obs_dim["policy"][0]
|
||||
self.history_length = history_length
|
||||
self.num_obs = self.num_one_step_obs * (self.history_length + 1)
|
||||
|
||||
# History buffer for policy observations
|
||||
self.obs_history_buf = torch.zeros(self.num_envs, self.num_obs, device=self.device)
|
||||
|
||||
# Termination tracking
|
||||
self._termination_ids = torch.tensor([], dtype=torch.long, device=self.device)
|
||||
|
||||
# Privileged observation setup if critic group exists
|
||||
if "critic" in self.unwrapped.observation_manager.group_obs_dim:
|
||||
self.num_one_step_privileged_obs = (
|
||||
self.unwrapped.observation_manager.group_obs_dim["critic"][0]
|
||||
)
|
||||
self.privileged_history_length = privileged_history_length
|
||||
self.num_privileged_obs = self.num_one_step_privileged_obs * (
|
||||
self.privileged_history_length + 1
|
||||
)
|
||||
self.privileged_obs_history_buf = torch.zeros(
|
||||
self.num_envs, self.num_privileged_obs, device=self.device
|
||||
)
|
||||
self._termination_privileged_obs = torch.zeros(
|
||||
0, self.num_privileged_obs, device=self.device
|
||||
)
|
||||
else:
|
||||
self.num_one_step_privileged_obs = None
|
||||
self.num_privileged_obs = None
|
||||
self.privileged_obs_history_buf = None
|
||||
self._termination_privileged_obs = None
|
||||
|
||||
# Reset at start since HIM runner does not call reset
|
||||
self.env.reset()
|
||||
|
||||
@property
|
||||
def cfg(self):
|
||||
return self.env.cfg
|
||||
|
||||
@property
|
||||
def render_mode(self):
|
||||
return self.env.render_mode
|
||||
|
||||
@property
|
||||
def observation_space(self):
|
||||
return self.env.observation_space
|
||||
|
||||
@property
|
||||
def action_space(self):
|
||||
return self.env.action_space
|
||||
|
||||
@classmethod
|
||||
def class_name(cls) -> str:
|
||||
return cls.__name__
|
||||
|
||||
@property
|
||||
def unwrapped(self):
|
||||
return self.env.unwrapped
|
||||
|
||||
@property
|
||||
def episode_length_buf(self) -> torch.Tensor:
|
||||
return self.unwrapped.episode_length_buf
|
||||
|
||||
@episode_length_buf.setter
|
||||
def episode_length_buf(self, value: torch.Tensor):
|
||||
self.unwrapped.episode_length_buf = value
|
||||
|
||||
def seed(self, seed: int = -1) -> int:
|
||||
return self.env.seed(seed)
|
||||
|
||||
def reset(self):
|
||||
obs_dict, _ = self.env.reset()
|
||||
policy_obs = obs_dict.get("policy", obs_dict[next(iter(obs_dict))])
|
||||
return policy_obs, {"observations": obs_dict}
|
||||
|
||||
def get_observations(self) -> torch.Tensor:
|
||||
return self.obs_history_buf
|
||||
|
||||
def get_privileged_observations(self) -> torch.Tensor | None:
|
||||
return self.privileged_obs_history_buf
|
||||
|
||||
def compute_termination_observations(
|
||||
self, env_ids: torch.Tensor, obs_before_reset: torch.Tensor
|
||||
) -> torch.Tensor | None:
|
||||
if len(env_ids) == 0:
|
||||
return torch.zeros(
|
||||
0, self.num_one_step_privileged_obs, device=self.device
|
||||
)
|
||||
return obs_before_reset[env_ids]
|
||||
|
||||
def step(self, actions: torch.Tensor):
|
||||
"""Execute one time-step and return 7 values for HIMLoco.
|
||||
|
||||
Returns:
|
||||
obs: History-stacked policy observations
|
||||
privileged_obs: Privileged/critic observations
|
||||
rewards: Rewards
|
||||
dones: Done flags (terminated | truncated)
|
||||
infos: Additional info dict
|
||||
termination_ids: Indices of envs that terminated this step
|
||||
termination_privileged_obs: Pre-reset privileged obs for terminated envs
|
||||
"""
|
||||
# Reorder actions from policy (Leg-by-Leg) to environment (Group)
|
||||
actions_reordered = actions[:, self.reorder_indices]
|
||||
|
||||
# Step the environment (6-value return)
|
||||
obs_dict, obs_before_reset, rewards, terminated, truncated, infos = \
|
||||
self.env.step(actions_reordered)
|
||||
|
||||
dones = (terminated | truncated).to(dtype=torch.long)
|
||||
|
||||
if not self.unwrapped.cfg.is_finite_horizon:
|
||||
infos["time_outs"] = truncated
|
||||
|
||||
# Extract policy observations from dict
|
||||
current_obs = obs_dict.get("policy")
|
||||
if current_obs is None:
|
||||
first_key = next(iter(obs_dict.keys()))
|
||||
current_obs = obs_dict[first_key]
|
||||
|
||||
# Update policy observation history buffer
|
||||
if self.history_length > 0:
|
||||
self.obs_history_buf = torch.cat(
|
||||
(
|
||||
current_obs[:, :self.num_one_step_obs],
|
||||
self.obs_history_buf[:, :-self.num_one_step_obs],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
else:
|
||||
self.obs_history_buf = current_obs
|
||||
|
||||
# Track terminated environments
|
||||
self._termination_ids = torch.nonzero(dones, as_tuple=False).squeeze(-1)
|
||||
|
||||
# Update privileged observation history buffer
|
||||
if "critic" in obs_dict and self.privileged_obs_history_buf is not None:
|
||||
current_privileged_obs = obs_dict["critic"]
|
||||
termination_observation = obs_before_reset.get("critic")
|
||||
if termination_observation is None:
|
||||
termination_observation = obs_before_reset.get("policy")
|
||||
self._termination_privileged_obs = self.compute_termination_observations(
|
||||
self._termination_ids, termination_observation
|
||||
)
|
||||
if self.privileged_history_length > 0:
|
||||
self.privileged_obs_history_buf = torch.cat(
|
||||
(
|
||||
current_privileged_obs[:, :self.num_one_step_privileged_obs],
|
||||
self.privileged_obs_history_buf[
|
||||
:, :-self.num_one_step_privileged_obs
|
||||
],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
else:
|
||||
self.privileged_obs_history_buf = current_privileged_obs
|
||||
|
||||
# NaN/Inf guard
|
||||
if torch.isnan(self.obs_history_buf).any() or torch.isinf(self.obs_history_buf).any():
|
||||
raise ValueError("NaN/Inf detected in obs_history_buf!")
|
||||
if torch.isnan(rewards).any() or torch.isinf(rewards).any():
|
||||
raise ValueError("NaN/Inf detected in rewards!")
|
||||
|
||||
return (
|
||||
self.obs_history_buf,
|
||||
self.privileged_obs_history_buf,
|
||||
rewards,
|
||||
dones,
|
||||
infos,
|
||||
self._termination_ids,
|
||||
self._termination_privileged_obs,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
return self.env.close()
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.env, name)
|
||||
@@ -0,0 +1,3 @@
|
||||
from .competition_terrains import RCWallTerrainCfg, RCLowBarTerrainCfg, RCPyramidStairsTerrainCfg
|
||||
|
||||
__all__ = ["RCWallTerrainCfg", "RCLowBarTerrainCfg", "RCPyramidStairsTerrainCfg"]
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Competition obstacle terrains for locomotion tasks.
|
||||
|
||||
Implemented in accordance with the interface specifications in
|
||||
mjlab/src/mjlab/terrains/primitive_terrains.py. Terrain difficulty (difficulty in [0, 1])
|
||||
is automatically supplied via TerrainGeneratorCfg and curriculum.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mjlab.terrains.terrain_generator import SubTerrainCfg, TerrainGeometry, TerrainOutput
|
||||
from mjlab.terrains.utils import make_plane
|
||||
from mjlab.terrains import BoxPyramidStairsTerrainCfg
|
||||
from mjlab.utils.color import brand_ramp
|
||||
|
||||
# Color constants matching primitive_terrains.py style
|
||||
_COLOR_ORANGE = (0.95, 0.55, 0.10)
|
||||
_COLOR_PURPLE = (0.60, 0.20, 0.80)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class RCWallTerrainCfg(SubTerrainCfg):
|
||||
"""Transverse wall obstacle terrain representing the race high wall.
|
||||
|
||||
The robot must sprint from the flat platform, vault over the wall, and proceed.
|
||||
As difficulty scales from 0 to 1, the wall height increases linearly from
|
||||
wall_height_range[0] to wall_height_range[1].
|
||||
|
||||
The actual physical wall specs are 1.0m width, 0.3m height, and 0.05m thickness.
|
||||
Difficulty range goes up to 0.35m to slightly exceed actual competition difficulty.
|
||||
"""
|
||||
|
||||
wall_height_range: tuple[float, float] = (0.0, 0.35)
|
||||
"""Wall height range (m), linearly interpolated with difficulty."""
|
||||
wall_thickness: float = 0.05
|
||||
"""Wall thickness (m)."""
|
||||
wall_length_frac: float = 0.8
|
||||
"""Wall length fraction of the terrain width (leaving gaps for visualization/debugging)."""
|
||||
platform_width: float = 1.5
|
||||
"""Sprint platform width (m)."""
|
||||
|
||||
def function(
|
||||
self,
|
||||
difficulty: float,
|
||||
spec: mujoco.MjSpec,
|
||||
rng: np.random.Generator,
|
||||
) -> TerrainOutput:
|
||||
del rng
|
||||
body = spec.body("terrain")
|
||||
geometries: list[TerrainGeometry] = []
|
||||
|
||||
wall_height = self.wall_height_range[0] + difficulty * (
|
||||
self.wall_height_range[1] - self.wall_height_range[0]
|
||||
)
|
||||
|
||||
# -- Ground floor base plane --
|
||||
floor_boxes = make_plane(body, self.size, 0.0, center_zero=False)
|
||||
floor_color = (0.45, 0.45, 0.45, 1.0)
|
||||
for box in floor_boxes:
|
||||
geometries.append(TerrainGeometry(geom=box, color=floor_color))
|
||||
|
||||
if wall_height < 1e-3:
|
||||
origin = np.array([self.size[0] / 2, self.size[1] / 2, 0.0])
|
||||
return TerrainOutput(origin=origin, geometries=geometries)
|
||||
|
||||
# -- Wall geometry: centered on x-midline, oriented along y-axis --
|
||||
wall_length = self.wall_length_frac * self.size[1]
|
||||
cx = self.size[0] / 2
|
||||
cy = self.size[1] / 2
|
||||
|
||||
wall_color = brand_ramp(_COLOR_ORANGE, difficulty)
|
||||
|
||||
wall_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX,
|
||||
size=(
|
||||
self.wall_thickness / 2, # half-size x
|
||||
wall_length / 2, # half-size y
|
||||
wall_height / 2, # half-size z
|
||||
),
|
||||
pos=(cx, cy, wall_height / 2),
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=wall_geom, color=wall_color))
|
||||
|
||||
# Spawn origin is set to the left platform area (e.g., 1.5m) to allow
|
||||
# the robot ample room to sprint/accelerate instead of spawning directly inside the obstacle.
|
||||
origin = np.array([1.5, cy, 0.0])
|
||||
return TerrainOutput(origin=origin, geometries=geometries)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class RCLowBarTerrainCfg(SubTerrainCfg):
|
||||
"""Low clearance bar obstacle terrain.
|
||||
|
||||
A horizontal bar is placed along the y-axis, requiring the robot to crouch and duck under.
|
||||
As difficulty scales from 0 to 1, the clearance height decreases linearly from
|
||||
clearance_range[1] to clearance_range[0] (harder is lower).
|
||||
|
||||
The actual competition spec is 0.30m離地 and 1.0m width.
|
||||
Here the range is configured as 0.25m to 0.35m to cover and exceed the standard specs.
|
||||
"""
|
||||
|
||||
clearance_range: tuple[float, float] = (0.25, 0.35)
|
||||
"""Bar ground clearance range (m). difficulty=0 -> highest/easiest, difficulty=1 -> lowest/hardest."""
|
||||
bar_radius: float = 0.025
|
||||
"""Bar cross-section radius (m)."""
|
||||
bar_length_frac: float = 0.85
|
||||
"""Bar length fraction of the terrain width (leaving gaps for visualization/debugging)."""
|
||||
|
||||
def function(
|
||||
self,
|
||||
difficulty: float,
|
||||
spec: mujoco.MjSpec,
|
||||
rng: np.random.Generator,
|
||||
) -> TerrainOutput:
|
||||
del rng
|
||||
body = spec.body("terrain")
|
||||
geometries: list[TerrainGeometry] = []
|
||||
|
||||
clearance = self.clearance_range[1] - difficulty * (
|
||||
self.clearance_range[1] - self.clearance_range[0]
|
||||
)
|
||||
bar_z = clearance + self.bar_radius # Cylinder center height
|
||||
|
||||
# -- Ground floor --
|
||||
floor_boxes = make_plane(body, self.size, 0.0, center_zero=False)
|
||||
floor_color = (0.45, 0.45, 0.45, 1.0)
|
||||
for box in floor_boxes:
|
||||
geometries.append(TerrainGeometry(geom=box, color=floor_color))
|
||||
|
||||
# -- Low bar cylinder (oriented along y-axis) --
|
||||
bar_length = self.bar_length_frac * self.size[1]
|
||||
cx = self.size[0] / 2
|
||||
cy = self.size[1] / 2
|
||||
|
||||
bar_color = brand_ramp(_COLOR_PURPLE, difficulty)
|
||||
|
||||
# MuJoCo cylinder defaults to z-axis orientation. Rotate 90 degrees around x-axis
|
||||
# to align with y-axis: quat [cos(pi/4), sin(pi/4), 0, 0] -> [0.7071, 0.7071, 0.0, 0.0]
|
||||
bar_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_CYLINDER,
|
||||
size=(self.bar_radius, bar_length / 2),
|
||||
pos=(cx, cy, bar_z),
|
||||
)
|
||||
bar_geom.quat = np.array([
|
||||
np.cos(np.pi / 4), np.sin(np.pi / 4), 0.0, 0.0
|
||||
])
|
||||
geometries.append(TerrainGeometry(geom=bar_geom, color=bar_color))
|
||||
|
||||
# -- Side supporting posts --
|
||||
post_color = bar_color
|
||||
post_half_h = bar_z / 2
|
||||
for y_off in [-bar_length / 2, bar_length / 2]:
|
||||
post = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_CYLINDER,
|
||||
size=(self.bar_radius, post_half_h),
|
||||
pos=(cx, cy + y_off, post_half_h),
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=post, color=post_color))
|
||||
|
||||
# Spawn origin is set to the left platform area (e.g., 1.5m) to allow
|
||||
# the robot ample room to sprint/accelerate instead of spawning directly under the bar.
|
||||
origin = np.array([1.5, cy, 0.0])
|
||||
return TerrainOutput(origin=origin, geometries=geometries)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class RCPyramidStairsTerrainCfg(BoxPyramidStairsTerrainCfg):
|
||||
"""Refactored pyramid stairs terrain with the spawn origin shifted to the left edge.
|
||||
|
||||
This ensures the robot can sprint/accelerate on flat ground before ascending,
|
||||
substantially reducing early falls or flips on vertical steps.
|
||||
"""
|
||||
|
||||
def function(
|
||||
self,
|
||||
difficulty: float,
|
||||
spec: mujoco.MjSpec,
|
||||
rng: np.random.Generator,
|
||||
) -> TerrainOutput:
|
||||
output = super().function(difficulty, spec, rng)
|
||||
cy = self.size[1] / 2
|
||||
# Relocate the spawn origin to the left platform zone
|
||||
output.origin = np.array([1.5, cy, 0.0])
|
||||
return output
|
||||
Reference in New Issue
Block a user