Robocon/RC_WheelLeg
4
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[train] 更新新MJCF与第一版完整训练框架

Browse Source
This commit is contained in:
eastzio
2026-07-27 12:31:21 +08:00
parent dfd646cf47
commit b08956aec7
116 changed files with 7649 additions and 6760 deletions
Download Patch File Download Diff File Expand all files Collapse all files
05_software/train/rc_mjlab/src/robot/config/env_cfgs.py
+64 -87
View File
@@ -74,7 +74,11 @@ from ..mdp.rewards import (
joint_mirror,
feet_contact_without_cmd,
upright_roll_only,
pitch_control_penalty,
upward,
joint_power,
ang_vel_xy_l2,
undesired_contacts,
contact_forces,
)
from ..mdp.curriculums import terrain_levels_vel_strict
from ..mdp.commands import UniformThresholdVelocityCommandCfg
@@ -174,7 +178,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
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),
"actions": ObservationTermCfg(func=velocity_mdp.last_action),
}
critic_terms = {
@@ -192,7 +196,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
observations = {
"actor": ObservationGroupCfg(
terms=actor_terms, concatenate_terms=True,
enable_corruption=True, history_length=6,
enable_corruption=True,
),
"critic": ObservationGroupCfg(
terms=critic_terms, concatenate_terms=True, enable_corruption=False,
@@ -239,8 +243,15 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
"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)},
"pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-math.pi, math.pi)},
"velocity_range": {
"x": (-0.5, 0.5),
"y": (-0.5, 0.5),
"z": (-0.5, 0.5),
"roll": (-0.5, 0.5),
"pitch": (-0.5, 0.5),
"yaw": (-0.5, 0.5),
},
"asset_cfg": SceneEntityCfg("robot"),
},
),
@@ -254,13 +265,9 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
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)},
params={"asset_cfg": SceneEntityCfg("robot", geom_names=(".*",)), "operation": "abs", "ranges": (0.3, 1.0)},
),
"actuator_stiffness": EventTermCfg(
func=envs_dr.joint_stiffness, mode="startup",
@@ -270,16 +277,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
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(
"body_mass_base": EventTermCfg(
func=envs_dr.body_mass, mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
@@ -287,16 +285,6 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
"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,
},
),
}
# ------------------
@@ -346,8 +334,8 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
),
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")),
metrics=metrics, curriculum=curriculum, decimation=4, episode_length_s=20.0,
sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic")),
viewer=ViewerConfig(body_name="base_link", distance=3.0, elevation=-20.0, azimuth=45.0),
)
@@ -380,33 +368,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
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)),
"pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.05, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
"pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.45, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
"random_grid": BoxRandomGridTerrainCfg(proportion=0.27, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
"random_rough": HfRandomUniformTerrainCfg(proportion=0.01, 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.01, 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.15, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
"sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.01, 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,
max_init_terrain_level=5,
)
# Disable default velocity stages command and bind strict velocity terrain curriculum
# Keep the custom terrain set, but align command/curriculum behavior with go2w rough.
cfg.curriculum.pop("command_vel", None)
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=terrain_levels_vel_strict, params={"command_name": "twist"})
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=velocity_mdp.terrain_levels_vel, 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"].rel_heading_envs = 1.0
cfg.commands["twist"].heading_control_stiffness = 0.5
cfg.commands["twist"].ranges.heading = (-math.pi, math.pi)
cfg.commands["twist"].rel_standing_envs = 0.2
cfg.commands["twist"].rel_standing_envs = 0.02
cfg.commands["twist"].ranges.lin_vel_x = (-1.0, 1.0)
cfg.commands["twist"].ranges.lin_vel_y = (-0.6, 0.6)
cfg.commands["twist"].ranges.ang_vel_z = (-1.0, 1.0)
# ------------------
# 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.pop("joint_friction", None)
cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.0), "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",
@@ -426,42 +417,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
# Rewards Integration
# ------------------
cfg.rewards["track_lin_vel"] = RewardTermCfg(
func=track_linear_velocity_l1,
weight=4.5,
func=track_linear_velocity,
weight=3.0,
params={"std": 0.5, "command_name": "twist"}
)
cfg.rewards["track_ang_vel"] = RewardTermCfg(
func=track_angular_velocity,
weight=1.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["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-2.0)
cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=ang_vel_xy_l2, weight=-0.05, params={"asset_cfg": SceneEntityCfg("robot")})
cfg.rewards.pop("upright", None)
cfg.rewards["roll_penalty"] = RewardTermCfg(
func=upright_roll_only,
weight=-1.0,
params={"asset_cfg": SceneEntityCfg("robot")}
)
cfg.rewards.pop("roll_penalty", None)
# 🌟 限制俯仰角死区（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.pop("action_rate_curriculum", None)
cfg.rewards["action_rate"].weight = -0.01
cfg.rewards["joint_torques"].weight = -1e-4
cfg.rewards["joint_torques"].weight = -2.5e-5
cfg.rewards["joint_power"] = RewardTermCfg(func=joint_power, weight=-2.0e-5)
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["joint_pos_limits"].weight = -5.0
cfg.rewards.pop("leg_motion_penalty", None)
cfg.rewards["is_terminated"].weight = 0.0
cfg.rewards.pop("leg_symmetry", None)
@@ -487,7 +472,7 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
cfg.rewards.pop("joint_deviation_l2", None)
cfg.rewards["joint_pos_penalty"] = RewardTermCfg(
func=joint_pos_penalty,
weight=-0.8,
weight=-1.0,
params={
"stand_still_scale": 5.0,
"velocity_threshold": 0.5,
@@ -502,37 +487,29 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
weight=0.1,
params={"command_name": "twist", "sensor_name": "feet_ground_contact"}
)
cfg.rewards["feet_air_time"].weight = 0.0
cfg.rewards["upward"] = RewardTermCfg(func=upward, weight=1.0)
cfg.rewards["base_height_l2"].weight = -0.5
cfg.rewards["base_height_l2"].weight = 0.0
cfg.rewards["base_height_l2"].params["target_height"] = 0.40
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["undesired_contacts"] = RewardTermCfg(func=undesired_contacts, weight=-1.0, params={"sensor_name": "body_collision", "threshold": 1.0})
cfg.rewards["contact_forces"] = RewardTermCfg(func=contact_forces, weight=-1.5e-4, params={"sensor_name": "feet_ground_contact", "threshold": 100.0})
# 🌟 严厉惩罚机身/胸部碰撞（防止硬撞高墙），逼迫机器人学会用前轮触墙并主动抬腿攀爬的“触觉反射”
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"):
for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel", "terrain_level_bonus", "action_rate_curriculum"):
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))
cfg.episode_length_s = 20.0
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, 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
@@ -663,7 +640,7 @@ def crawl_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
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))
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, 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)
05_software/train/rc_mjlab/src/robot/config/rl_cfg.py
+4 -6
View File
@@ -33,10 +33,10 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
value_loss_coef=1.0,
use_clipped_value_loss=True,
clip_param=0.2,
entropy_coef=0.002,
entropy_coef=0.001,#第一轮为0.003
num_learning_epochs=5,
num_mini_batches=4,
learning_rate=2.0e-4,
learning_rate=8.0e-4,
schedule="adaptive",
gamma=0.99,
lam=0.95,
@@ -44,9 +44,9 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
max_grad_norm=1.0,
),
experiment_name="robot_rough",
save_interval=50,
save_interval=100,
num_steps_per_env=24,
max_iterations=15_000,
max_iterations=20_000,
)
@@ -55,5 +55,3 @@ def crawl_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
cfg.experiment_name = "robot_crawl"
return cfg
05_software/train/rc_mjlab/src/robot/mdp/rewards.py
+73
View File
@@ -28,6 +28,7 @@ def track_linear_velocity(
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)
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
@@ -45,6 +46,7 @@ def track_angular_velocity(
actual = asset.data.root_link_ang_vel_b
z_error = torch.square(command[:, 2] - actual[:, 2])
reward = torch.exp(-z_error / std**2)
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
@@ -94,6 +96,7 @@ def base_height_l2(
error = root_z - target_height
reward = torch.square(error)
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
@@ -304,6 +307,7 @@ def stand_still(
angular_norm = torch.abs(command[:, 2])
inactive = (linear_norm + angular_norm < command_threshold).float()
reward = cost * inactive
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def hip_deviation(
@@ -355,6 +359,20 @@ def lin_vel_z_l2(
asset_cfg = SceneEntityCfg("robot")
asset: Entity = env.scene[asset_cfg.name]
reward = torch.square(asset.data.root_link_lin_vel_b[:, 2])
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def ang_vel_xy_l2(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg | None = None,
) -> torch.Tensor:
"""Penalize xy-axis base angular velocity using the go2w kernel."""
if asset_cfg is None:
asset_cfg = SceneEntityCfg("robot")
asset: Entity = env.scene[asset_cfg.name]
reward = torch.sum(torch.square(asset.data.root_link_ang_vel_b[:, :2]), dim=1)
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
@@ -650,6 +668,8 @@ def feet_contact_without_cmd(env, command_name: str, sensor_name: str) -> torch.
linear_norm = torch.norm(cmd[:, :2], dim=1)
angular_norm = torch.abs(cmd[:, 2])
reward *= (linear_norm + angular_norm) < 0.1
asset = env.scene["robot"]
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def joint_pos_penalty(
@@ -674,6 +694,7 @@ def joint_pos_penalty(
running_reward,
stand_still_scale * running_reward,
)
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
@@ -694,6 +715,50 @@ def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg
)
reward += diff
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def undesired_contacts(
env: ManagerBasedRlEnv,
sensor_name: str,
threshold: float = 1.0,
) -> torch.Tensor:
"""Penalize non-wheel contacts above a force threshold."""
from mjlab.sensor import ContactSensor
sensor: ContactSensor = env.scene[sensor_name]
data = sensor.data
if data.force_history is not None:
force_mag = torch.norm(data.force_history, dim=-1)
is_contact = torch.max(force_mag, dim=2)[0] > threshold
else:
force_mag = torch.norm(data.force, dim=-1)
is_contact = force_mag > threshold
reward = torch.sum(is_contact, dim=1).float()
asset = env.scene["robot"]
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
def contact_forces(
env: ManagerBasedRlEnv,
sensor_name: str,
threshold: float = 100.0,
) -> torch.Tensor:
"""Penalize foot contact forces above threshold."""
from mjlab.sensor import ContactSensor
sensor: ContactSensor = env.scene[sensor_name]
data = sensor.data
if data.force_history is not None:
force_mag = torch.norm(data.force_history, dim=-1)
peak_force = torch.max(force_mag, dim=2)[0]
else:
peak_force = torch.norm(data.force, dim=-1)
reward = torch.sum(torch.clamp(peak_force - threshold, min=0.0), dim=1)
asset = env.scene["robot"]
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
return reward
@@ -705,6 +770,14 @@ def upward(env, asset_cfg=None):
reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
return reward
def joint_power(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
"""Penalty for total joint mechanical power: sum(|tau * dq|)."""
if asset_cfg is None:
asset_cfg = SceneEntityCfg("robot")
asset: Entity = env.scene[asset_cfg.name]
return torch.sum(torch.abs(asset.data.qfrc_actuator * asset.data.joint_vel), dim=1)
def upright_roll_only(env, asset_cfg=None):
if asset_cfg is None:
from mjlab.envs.manager_based_rl_env import SceneEntityCfg
05_software/train/rc_mjlab/src/robot/rl/himloco/__init__.py
-41
View File
@@ -1,41 +0,0 @@
# 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 *
05_software/train/rc_mjlab/src/robot/rl/himloco/algorithms/__init__.py
-31
View File
@@ -1,31 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/algorithms/him_ppo.py
-192
View File
@@ -1,192 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/config/__init__.py
-31
View File
@@ -1,31 +0,0 @@
# 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 *
05_software/train/rc_mjlab/src/robot/rl/himloco/config/rl_cfg.py
-180
View File
@@ -1,180 +0,0 @@
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."""
05_software/train/rc_mjlab/src/robot/rl/himloco/env/__init__.py
Vendored
-31
View File
@@ -1,31 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/env/vec_env.py
Vendored
-60
View File
@@ -1,60 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/modules/__init__.py
-32
View File
@@ -1,32 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/modules/him_actor_critic.py
-236
View File
@@ -1,236 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/modules/him_estimator.py
-155
View File
@@ -1,155 +0,0 @@
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
05_software/train/rc_mjlab/src/robot/rl/himloco/runners/__init__.py
-31
View File
@@ -1,31 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/runners/him_on_policy_runner.py
-362
View File
@@ -1,362 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/storage/__init__.py
-4
View File
@@ -1,4 +0,0 @@
# Copyright 2021 ETH Zurich, NVIDIA CORPORATION
# SPDX-License-Identifier: BSD-3-Clause
from .him_rollout_storage import HIMRolloutStorage
05_software/train/rc_mjlab/src/robot/rl/himloco/storage/him_rollout_storage.py
-167
View File
@@ -1,167 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/utils/__init__.py
-31
View File
@@ -1,31 +0,0 @@
# 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
05_software/train/rc_mjlab/src/robot/rl/himloco/utils/utils.py
-71
View File
@@ -1,71 +0,0 @@
# 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)
05_software/train/rc_mjlab/src/robot/rl/himloco/wrappers/__init__.py
View File
05_software/train/rc_mjlab/src/robot/robot_cfg.py
+6 -6
View File
@@ -27,13 +27,13 @@ def get_spec() -> mujoco.MjSpec:
# 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
# Leg joints: Position PD control aligned with current real deployment.
STIFFNESS_LEG = 50.0
DAMPING_LEG = 1.5
EFFORT_LEG = 17.0
# Wheel joints: Velocity control with damping = 0.5 and effort limit = 17 Nm.
DAMPING_WHEEL = 0.5
# Wheel joints: Velocity control aligned with current real deployment.
DAMPING_WHEEL = 1.0
EFFORT_WHEEL = 17.0
# Maximum joint speed for all actuators (rad/s)
@@ -110,4 +110,4 @@ def get_robot_crawl_cfg() -> EntityCfg:
# 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
WHEEL_VEL_SCALE = 5.0
05_software/train/rc_mjlab/src/robot/terrains/competition_terrains.py
+18 -16
View File
@@ -24,9 +24,9 @@ _COLOR_PURPLE = (0.60, 0.20, 0.80)
@dataclass(kw_only=True)
class RCWallTerrainCfg(SubTerrainCfg):
"""Transverse wall obstacle terrain representing the race high wall.
"""Triple transverse wall obstacle terrain representing repeated race high walls.
The robot must sprint from the flat platform, vault over the wall, and proceed.
The robot must sprint from the flat platform, vault over three walls, and proceed.
As difficulty scales from 0 to 1, the wall height increases linearly from
wall_height_range[0] to wall_height_range[1].
@@ -42,6 +42,8 @@ class RCWallTerrainCfg(SubTerrainCfg):
"""Wall length fraction of the terrain width (leaving gaps for visualization/debugging)."""
platform_width: float = 1.5
"""Sprint platform width (m)."""
wall_centers_x: tuple[float, float, float] = (2.9, 4.45, 6.0)
"""Wall center positions along x, spaced to keep a short sprint, two recovery gaps, and exit room."""
def function(
self,
@@ -67,26 +69,26 @@ class RCWallTerrainCfg(SubTerrainCfg):
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 geometry: three transverse walls 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))
for cx in self.wall_centers_x:
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.
# Spawn origin is set to the left platform area to allow a short sprint
# before the first wall and limited recovery space between subsequent walls.
origin = np.array([1.5, cy, 0.0])
return TerrainOutput(origin=origin, geometries=geometries)