From 3a68a35be7bb90ad4a94d21cf9ea84639c9fa7a1 Mon Sep 17 00:00:00 2001 From: wertyuilife Date: Wed, 8 Apr 2026 17:09:42 +0800 Subject: [PATCH] replace commands.py with commands_go2_rl_gym.py. --- .../robot_lab/tasks/go2/mdp/__init__.py | 1 - .../robot_lab/tasks/go2/mdp/commands.py | 673 +++++++++--------- .../tasks/go2/mdp/commands_go2_rl_gym.py | 422 ----------- 3 files changed, 349 insertions(+), 747 deletions(-) delete mode 100644 source/robot_lab/robot_lab/tasks/go2/mdp/commands_go2_rl_gym.py diff --git a/source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py b/source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py index 7f57ddb..bda76aa 100644 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py +++ b/source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py @@ -17,4 +17,3 @@ from .events import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .utils import * # noqa: F401, F403 -from .commands_go2_rl_gym import * # noqa: F401, F403 diff --git a/source/robot_lab/robot_lab/tasks/go2/mdp/commands.py b/source/robot_lab/robot_lab/tasks/go2/mdp/commands.py index b0b99fd..5f52e92 100644 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/commands.py +++ b/source/robot_lab/robot_lab/tasks/go2/mdp/commands.py @@ -1,343 +1,317 @@ -# Copyright (c) 2024-2025 Ziqi Fan -# SPDX-License-Identifier: Apache-2.0 +# -*- coding: utf-8 -*- +''' +@File : commands_go2_rl_gym.py +@Time : 2026/04/01 17:16:44 +@Author : wty-yy +@Version : 1.0 +@Blog : https://wty-yy.github.io/ +@Desc : CommandTerm for go2_rl_gym style command generation, reference to https://github.com/wty-yy/go2_rl_gym -from __future__ import annotations +IsaacLab CommandTerm working flow: +Env: after compute reward call command.compute(dt) +1. self._update_metrics(): update self.metrics dict for logging +2. self.time_left -= dt +3. self._resample(self.time_left <= 0) +4. self._update_command(): update command if needed + +Get command from self.command property, return command, shape=(num_envs, command_dim) + +Note: +1. We don't use original self._resample(env_ids) and self._resample_command(env_ids), because it will randomize time_left +2. Remove heading command +3. We don't use curriculum item to update curriculum, inplace update +''' + +from __future__ import annotations # For forward reference of type hints + +from typing import TYPE_CHECKING, Sequence +if TYPE_CHECKING: # Avoid circular import for type checking + from robot_lab.tasks.go2.env.go2_env import ActionDelayGo2Env -import time -import torch -import copy -from collections.abc import Sequence -from typing import TYPE_CHECKING -from dataclasses import MISSING from itertools import product -from isaaclab.managers import CommandTerm +import torch + from isaaclab.utils import configclass +from isaaclab.managers import CommandTerm, CommandTermCfg +from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg +from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG from isaaclab.assets import Articulation -from isaaclab.envs.mdp import UniformVelocityCommandCfg import isaaclab.utils.math as math_utils -from isaaclab.markers import VisualizationMarkers -# SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2022-2025, The Isaac Lab Project Developers - -from .utils import is_robot_on_terrain - -if TYPE_CHECKING: - from robot_lab.tasks.go2.env.go2_env import ManagerBasedRLEnv +from robot_lab.tasks.go2.mdp.utils import is_robot_on_terrain, sample_disjoint_intervals, sample_single_interval -class UniformVelTerrainCmd(CommandTerm): - cfg: UniformVelTerrainCmdCfg +class Go2RLGymCommand(CommandTerm): + cfg: Go2RLGymCommandCfg + _env: ActionDelayGo2Env - def __init__(self, cfg: UniformVelTerrainCmdCfg, env: ManagerBasedRLEnv): + def __init__(self, cfg: Go2RLGymCommandCfg, env: ActionDelayGo2Env): + """Reference: https://github.com/wty-yy/go2_rl_gym/blob/master/legged_gym/envs/base/legged_robot.py + LeggedRobot._resample_command() and LeggedRobot._post_physics_step_callback() + """ super().__init__(cfg, env) - # obtain the robot asset - # -- robot - self.robot: Articulation = env.scene[cfg.asset_name] - - # crete buffers to store the command - # -- command: x vel, y vel, yaw vel, heading - self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device) - self.heading_target = torch.zeros(self.num_envs, device=self.device) - self.is_heading_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - self.is_standing_env = torch.zeros_like(self.is_heading_env) - - # -- metrics - self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device) - self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device) - if self.cfg.use_phase_cmd: - self.phase = torch.zeros(self.num_envs, 1, device=self.device) - self.cycle_time = cfg.cycle_time - self.max_angular_envs = cfg.max_angular_envs - self.limit_vel_envs = cfg.limit_vel_envs - self.stop_heading = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - - # 获取当前环境的地形类型索引 - self.terrain_types = list(env.scene.terrain.cfg.terrain_generator.sub_terrains.keys()) - if not isinstance(self.cfg.ranges, dict): - single_range = cfg.ranges - expanded_ranges = {} - for t_name in self.terrain_types: - expanded_ranges[t_name] = copy.deepcopy(single_range) - self.cfg.ranges = expanded_ranges - assert set(self.terrain_types) == set(list(self.cfg.ranges.keys())), \ - "Terrain types in cfg.ranges do not match those in terrain generator config." - self.terrain_type_to_id = {name: i for i, name in enumerate(self.terrain_types)} - self.global_env_terrain_idx = torch.full((self.num_envs,), -1, dtype=torch.long, device=self.device) - for t_type in self.terrain_types: - ids = is_robot_on_terrain(self._env, t_type).nonzero(as_tuple=True)[0] - if len(ids) > 0: - self.global_env_terrain_idx[ids] = self.terrain_type_to_id[t_type] - # 初始化累计指令和最大位移 - self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, device=self.device) - self.dt = self._env.step_dt - self.max_episode_length = self._env.max_episode_length - self.target_dist = env.scene.terrain.cfg.terrain_generator.size[0] * 0.625 - self.max_move_distance = torch.zeros(self.num_envs, device=self.device) - self.env_origins = env.scene.env_origins - - # 极限指令组合 - limit_options_x = [-1, 1] - limit_options_y = [-1, 1] - limit_options_z = [-1, 0, 1] # 允许 Yaw 为 0，即直线冲刺 - self.limit_vel_combinations = torch.tensor( - list(product(limit_options_x, limit_options_y, limit_options_z)), - dtype=torch.long, - device=self.device - ) - + self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, dtype=torch.float, device=self.device) + self.max_move_distance = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) self.last_is_limit_vel = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self.commands = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device) # [lin_vel_x, lin_vel_y, ang_vel_yaw] + self.command_ranges = self.cfg.ranges.to_dict() + self.env_command_ranges = { + 'lin_vel_x': torch.tensor(self.command_ranges['lin_vel_x'], device=self.device).repeat(self.num_envs, 1), + 'lin_vel_y': torch.tensor(self.command_ranges['lin_vel_y'], device=self.device).repeat(self.num_envs, 1), + 'ang_vel_yaw': torch.tensor(self.command_ranges['ang_vel_yaw'], device=self.device).repeat(self.num_envs, 1), + } + self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), + abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) + self.limit_vel_comb = torch.tensor(list(product( + self.cfg.limit_vel["lin_vel_x"], + self.cfg.limit_vel["lin_vel_y"], + self.cfg.limit_vel["ang_vel_yaw"] + )), device=self.device) + self._init_terrain_infos() + self._update_env_command_ranges() + self.robot: Articulation = env.scene[cfg.asset_name] + self.zero_command_prob = 0 + self.max_command_x = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) + + self.cfg.command_range_curriculum = sorted(self.cfg.command_range_curriculum, key=lambda x: x['iter'], reverse=True) def __str__(self) -> str: - """Return a string representation of the command generator.""" - msg = "UniformVelocityCommand:\n" - msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" - msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" - msg += f"\tHeading command: {self.cfg.heading_command}\n" - if self.cfg.heading_command: - msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n" - msg += f"\tStanding probability: {self.cfg.rel_standing_envs}\n" - msg += f"\tMax angular vel envs: {self.max_angular_envs}\n" - msg += f"\tLimit linear vel envs: {self.limit_vel_envs}\n" - for t_type in self.terrain_types: - msg += f"\tTerrain type '{t_type}' command ranges: {self.cfg.ranges[t_type]}" + """Return a string representation of the command term.""" + msg = (f"""Go2RLGymCommand:\n""" + f"""Command shape: {self.commands.shape}""") return msg + def _init_terrain_infos(self): + """Initialize terrain types and indices for each environment.""" + self.terrain_types = list(self._env.scene.terrain.cfg.terrain_generator.sub_terrains.keys()) + for terrain_type in self.terrain_types: + if terrain_type not in self.cfg.terrain_max_command_ranges: + raise ValueError(f"Terrain type '{terrain_type}' is not defined in cfg.terrain_max_command_ranges.") + self.terrain_type2idx = {terrain_type: idx for idx, terrain_type in enumerate(self.terrain_types)} + self.terrain_idxs = torch.full((self.num_envs,), -1, dtype=torch.long, device=self.device) + for terrain_type in self.terrain_types: + idxs = is_robot_on_terrain(self._env, terrain_type).nonzero().flatten() + if len(idxs) > 0: + self.terrain_idxs[idxs] = self.terrain_type2idx[terrain_type] + self.terrain_length = self._env.scene.terrain.cfg.terrain_generator.size[0] + @property def command(self) -> torch.Tensor: - if self.cfg.use_phase_cmd: - phase = self.phase * 2 * torch.pi - return torch.cat([self.vel_command_b, torch.sin(phase), torch.cos(phase)], dim=-1) - return self.vel_command_b + return self.commands def _update_metrics(self): - # time for which the command was executed - max_command_time = self.cfg.resampling_time_range[1] - max_command_step = max_command_time / self._env.step_dt - # logs data - self.metrics["error_vel_xy"] += ( - torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step - ) - self.metrics["error_vel_yaw"] += ( - torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step - ) - + self.max_command_x[:] = self.command_ranges["lin_vel_x"][1] + self.metrics["max_command_x"] = self.max_command_x + def reset(self, env_ids: Sequence[int] | None = None): + self.time_left[env_ids] = self.cfg.resampling_time self.commands_xy_accumulation[env_ids] = 0.0 self.max_move_distance[env_ids] = 0.0 - extra = super().reset(env_ids) - return extra + return super().reset(env_ids) + + def _resample(self, env_ids: Sequence[int]): + """ Randommly select commands of some environments - def _terrain_cmd_sample(self, dim: str, env_ids: torch.Tensor, min_abs_vel: torch.Tensor | None = None): - """根据不同地形sample不同指令""" - cmd = torch.zeros(len(env_ids), device=self.device) - batch_terrain_idxs = self.global_env_terrain_idx[env_ids] - - for t_type in self.terrain_types: - t_id = self.terrain_type_to_id[t_type] - mask = (batch_terrain_idxs == t_id) - - if mask.any(): - count = mask.sum().item() - if dim == "x": - r = self.cfg.ranges[t_type].lin_vel_x - elif dim == "y": - r = self.cfg.ranges[t_type].lin_vel_y - elif dim == "z": - r = self.cfg.ranges[t_type].ang_vel_z - elif dim == "heading": - r = self.cfg.ranges[t_type].heading - - sampled_vals = torch.empty(count, device=self.device).uniform_(*r) # type: ignore + Args: + env_ids (List[int]): Environments ids for which new commands are needed + """ + env = self._env + if len(env_ids) == 0: + return + # update command curriculum with train steps + if len(self.cfg.command_range_curriculum): + current_iter = env.common_step_counter // self.cfg.num_steps_per_iter + for i in range(len(self.cfg.command_range_curriculum)-1, -1, -1): # iterate backwards to be able to pop entries + cfg = self.cfg.command_range_curriculum[i] + if current_iter >= cfg["iter"]: + self.command_ranges["lin_vel_x"] = cfg["lin_vel_x"] + self.command_ranges["lin_vel_y"] = cfg["lin_vel_y"] + self.command_ranges["ang_vel_yaw"] = cfg["ang_vel_yaw"] + self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), + abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) + self.cfg.command_range_curriculum.pop(i) + self._update_env_command_ranges() + print(f"Command range updated at iter {current_iter}: {self.command_ranges}") + remaining_dist = torch.clip(0.625 * self.terrain_length - torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.resampling_time, 0.0) + self.time_left[env_ids] = self.cfg.resampling_time + if self.cfg.dynamic_resample_commands: + # arrive at boundary 0.625 times the width of the remaining distance + if ((env.max_episode_length - env.episode_length_buf[env_ids]) + 1 == 0).any(): + raise ValueError("Some envs have zero remaining episode length during command resampling") + vel_low_bound = torch.clip(remaining_dist / ((env.max_episode_length - env.episode_length_buf[env_ids] + 1 + 1e-9) * env.step_dt), 0.0) + self.commands[env_ids, 0] = sample_disjoint_intervals( + env_ids, + vel_low_bound, + self.env_command_ranges["lin_vel_x"][env_ids, 0], + self.env_command_ranges["lin_vel_x"][env_ids, 1], + self.device + ) + self.commands[env_ids, 1] = sample_disjoint_intervals( + env_ids, + vel_low_bound, + self.env_command_ranges["lin_vel_y"][env_ids, 0], + self.env_command_ranges["lin_vel_y"][env_ids, 1], + self.device + ) + r = torch.rand(len(env_ids), device=self.device) + lower = self.env_command_ranges["ang_vel_yaw"][env_ids, 0] + upper = self.env_command_ranges["ang_vel_yaw"][env_ids, 1] + self.commands[env_ids, 2] = (upper - lower) * r + lower + else: + self.commands[env_ids, 0] = sample_single_interval( + env_ids, + self.env_command_ranges["lin_vel_x"][env_ids, 0], + self.env_command_ranges["lin_vel_x"][env_ids, 1], + self.device + ) + self.commands[env_ids, 1] = sample_single_interval( + env_ids, + self.env_command_ranges["lin_vel_y"][env_ids, 0], + self.env_command_ranges["lin_vel_y"][env_ids, 1], + self.device + ) + self.commands[env_ids, 2] = sample_single_interval( + env_ids, + self.env_command_ranges["ang_vel_yaw"][env_ids, 0], + self.env_command_ranges["ang_vel_yaw"][env_ids, 1], + self.device + ) - if min_abs_vel is not None and (dim == "x" or dim == "y"): - lower_bound = min_abs_vel[mask] - - max_abs = max(abs(r[0]), abs(r[1])) - lower_bound = torch.clamp(lower_bound, max=max_abs) + # set small commands to zero + self.commands[env_ids, :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) - # 重新采样 Magnitude: [lower_bound, max_abs] - mag = torch.empty(count, device=self.device).uniform_(0, 1) * (max_abs - lower_bound) + lower_bound - - if r[0] >= 0: - sampled_vals = mag - elif r[1] <= 0: - sampled_vals = -mag - else: - sign = torch.sign(torch.empty(count, device=self.device).uniform_(-1, 1)) - sign = torch.where(sign == 0, torch.ones_like(sign), sign) - sampled_vals = mag * sign - - cmd[mask] = sampled_vals - - return cmd - - def _get_cmd_board(self, dim: str, env_ids: torch.Tensor): - """根据预计算的地形类型获取速度指令边界""" - cmd_min = torch.zeros(len(env_ids), device=self.device) - cmd_max = torch.zeros(len(env_ids), device=self.device) - - batch_terrain_idxs = self.global_env_terrain_idx[env_ids] - - for t_type in self.terrain_types: - t_id = self.terrain_type_to_id[t_type] - mask = (batch_terrain_idxs == t_id) - - if mask.any(): - if dim == "x": - r = self.cfg.ranges[t_type].lin_vel_x - elif dim == "y": - r = self.cfg.ranges[t_type].lin_vel_y - elif dim == "z": - r = self.cfg.ranges[t_type].ang_vel_z - else: - continue - - cmd_min[mask] = r[0] - cmd_max[mask] = r[1] - - return cmd_min, cmd_max - - def _resample_command(self, env_ids: Sequence[int]): - _env_ids = torch.tensor(env_ids, device=self.device) - - # 计算剩余距离: 目标距离 - 已经指令累积走过的距离 * 上一次重采样时间 - dist_covered = torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.resampling_time_range[0] - remaining_dist = torch.clamp(self.target_dist - dist_covered, min=0.0) - - # 计算剩余时间: (最大步数 - 当前步数) * dt - time_left = (self.max_episode_length - self._env.episode_length_buf[env_ids]) * self.dt - # 计算下限速度: 距离 / 时间 - vel_low_bound = torch.zeros(len(env_ids), device=self.device) - # 避免除以零或负数时间 - valid_time_mask = time_left > 1e-4 - if valid_time_mask.any(): - vel_low_bound[valid_time_mask] = remaining_dist[valid_time_mask] / time_left[valid_time_mask] - - # 根据不同地形类型采样速度指令 - self.vel_command_b[env_ids, 0] = self._terrain_cmd_sample("x", _env_ids, min_abs_vel=vel_low_bound) - self.vel_command_b[env_ids, 0] = torch.where( - self.vel_command_b[env_ids, 0].abs() < 0.1, - torch.zeros_like(self.vel_command_b[env_ids, 0]), - self.vel_command_b[env_ids, 0] + rand_prob = torch.rand(len(env_ids), device=self.device) + min_prob, max_prob = 0.0, 0.0 + # set limitation lin vel + if self.cfg.limit_vel_prob > 0.0: + max_prob += self.cfg.limit_vel_prob + lim_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) + lim_env_ids = env_ids[lim_mask] + if len(lim_env_ids) > 0: + change_lim_env_ids = lim_env_ids + if self.cfg.limit_vel_invert_when_continuous: + was_limited = self.last_is_limit_vel[lim_env_ids] + invert_env_ids = lim_env_ids[was_limited] + self.commands[invert_env_ids, 0] *= -1.0 + self.commands[invert_env_ids, 1] *= -1.0 + self.commands[invert_env_ids, 2] *= -1.0 + change_lim_env_ids = lim_env_ids[~was_limited] + vel_idx = torch.randint(0, self.limit_vel_comb.shape[0], (len(change_lim_env_ids),), device=self.device) + lin_vel_x_lim = torch.where( + self.limit_vel_comb[vel_idx, 0] == -1, + self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 0], + self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 1], ) - self.vel_command_b[env_ids, 1] = self._terrain_cmd_sample("y", _env_ids, min_abs_vel=vel_low_bound) - self.vel_command_b[env_ids, 1] = torch.where( - self.vel_command_b[env_ids, 1].abs() < 0.1, - torch.zeros_like(self.vel_command_b[env_ids, 1]), - self.vel_command_b[env_ids, 1] + lin_vel_x_lim[self.limit_vel_comb[vel_idx, 0] == 0] = 0.0 + lin_vel_y_lim = torch.where( + self.limit_vel_comb[vel_idx, 1] == -1, + self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 0], + self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 1] ) - self.vel_command_b[env_ids, 2] = self._terrain_cmd_sample("z", _env_ids) - self.vel_command_b[env_ids, 2] = torch.where( - self.vel_command_b[env_ids, 2].abs() < 0.1, - torch.zeros_like(self.vel_command_b[env_ids, 2]), - self.vel_command_b[env_ids, 2] + lin_vel_y_lim[self.limit_vel_comb[vel_idx, 1] == 0] = 0.0 + ang_vel_z_lim = torch.where( + self.limit_vel_comb[vel_idx, 2] == -1, + self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 0], + self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 1] ) - - # 计算heading和stand - r = torch.empty(len(env_ids), device=self.device) - if self.cfg.heading_command: - self.heading_target[env_ids] = self._terrain_cmd_sample("heading", _env_ids) - # update heading envs - self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs - # update standing envs - self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs - - # 静止环境sample最大转向 - standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() - if len(standing_env_ids) > 0: - self.vel_command_b[standing_env_ids, :] = 0.0 - ang_vel_rand = torch.rand(len(standing_env_ids), device=self.device) - add_ang_mask = ang_vel_rand < self.max_angular_envs - add_ang_env_ids = standing_env_ids[add_ang_mask] - if self.max_angular_envs > 0 and len(add_ang_env_ids) > 0: - direction_rand = torch.rand(len(add_ang_env_ids), device=self.device) - min_z, max_z = self._get_cmd_board("z", add_ang_env_ids) - self.vel_command_b[add_ang_env_ids, 2] = torch.where(direction_rand < 0.5, min_z, max_z) - self.stop_heading[add_ang_env_ids] = True - - limit_vel_env_ids = (self.is_standing_env == 0).nonzero(as_tuple=False).flatten() - - # 非静止环境sample极限速度 - current_ids = env_ids - mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - mask[current_ids] = True - target_limit_ids = limit_vel_env_ids[mask[limit_vel_env_ids]] + ang_vel_z_lim[self.limit_vel_comb[vel_idx, 2] == 0] = 0.0 + self.commands[change_lim_env_ids, 0] = lin_vel_x_lim + self.commands[change_lim_env_ids, 1] = lin_vel_y_lim + self.commands[change_lim_env_ids, 2] = ang_vel_z_lim + self.last_is_limit_vel[env_ids] = False + self.last_is_limit_vel[lim_env_ids] = True + else: + self.last_is_limit_vel[env_ids] = False + min_prob += self.cfg.limit_vel_prob - if self.limit_vel_envs > 0 and len(target_limit_ids) > 0: - limit_prob = torch.rand(len(target_limit_ids), device=self.device) - execute_limit_mask = limit_prob < self.limit_vel_envs - execute_ids = target_limit_ids[execute_limit_mask] + # set all commands to zero with some probability + if self.cfg.zero_command_curriculum is not None: + self.zero_command_prob = self.get_current_scale(self.cfg.zero_command_curriculum) + if self.zero_command_prob > 0.0: + max_prob += self.zero_command_prob + next_time_left = torch.clip( + env.max_episode_length_s - env.episode_length_buf[env_ids] * env.step_dt - (remaining_dist / (0.8 * self.max_lin_vel + 1e-9)), + min=0.0, + max=self.cfg.resampling_time, + ) + zero_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) * (next_time_left > 0.0) + zero_env_ids = env_ids[zero_mask] + if len(zero_env_ids) > 0: + self.commands[zero_env_ids, :2] = 0.0 + self.time_left[zero_env_ids] = next_time_left[zero_mask] + if self.cfg.limit_ang_vel_at_zero_command_prob > 0.0: + ang_vel_rand = torch.rand(len(zero_env_ids), device=self.device) # independent distribution + add_ang_mask = ang_vel_rand < self.cfg.limit_ang_vel_at_zero_command_prob + add_ang_env_ids = zero_env_ids[add_ang_mask] + if len(add_ang_env_ids) > 0: + direction_rand = torch.rand(len(add_ang_env_ids), device=self.device) + self.commands[add_ang_env_ids, 2] = torch.where( + direction_rand < 0.5, + self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 0], + self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 1] + ) + min_prob += self.zero_command_prob + + self.commands_xy_accumulation[env_ids] += self.commands[env_ids, :2] - if len(execute_ids) > 0: - - num_combs = self.limit_vel_combinations.shape[0] - comb_indices = torch.randint(0, num_combs, (len(execute_ids),), device=self.device) - selected_combs = self.limit_vel_combinations[comb_indices] # (N, 3) - - min_x, max_x = self._get_cmd_board("x", execute_ids) - min_y, max_y = self._get_cmd_board("y", execute_ids) - min_z, max_z = self._get_cmd_board("z", execute_ids) - - vals_x = torch.zeros_like(min_x) - vals_x = torch.where(selected_combs[:, 0] == -1, min_x, vals_x) - vals_x = torch.where(selected_combs[:, 0] == 1, max_x, vals_x) - - vals_y = torch.zeros_like(min_y) - vals_y = torch.where(selected_combs[:, 1] == -1, min_y, vals_y) - vals_y = torch.where(selected_combs[:, 1] == 1, max_y, vals_y) - - vals_z = torch.zeros_like(min_z) - vals_z = torch.where(selected_combs[:, 2] == -1, min_z, vals_z) - vals_z = torch.where(selected_combs[:, 2] == 1, max_z, vals_z) - - self.vel_command_b[execute_ids, 0] = vals_x - self.vel_command_b[execute_ids, 1] = vals_y - self.vel_command_b[execute_ids, 2] = vals_z - - self.stop_heading[execute_ids] = True - - if self.cfg.use_phase_cmd: - self.phase[env_ids] = torch.rand((len(env_ids), 1), device=self.device) - self.commands_xy_accumulation[env_ids] += self.vel_command_b[env_ids, :2] - def _update_command(self): - # 停止heading更新 - stop_heading_env_ids = self.stop_heading.nonzero(as_tuple=False).flatten() - if len(stop_heading_env_ids) > 0: - self.is_heading_env[stop_heading_env_ids] = False - - # Compute angular velocity from heading direction - if self.cfg.heading_command: - # resolve indices of heading envs - env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten() - # compute angular velocity - if len(env_ids) > 0: - heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids]) - min_z, max_z = self._get_cmd_board("z", env_ids) - - self.vel_command_b[env_ids, 2] = torch.clip( - self.cfg.heading_control_stiffness * heading_error, - min=min_z, - max=max_z, - ) - if self.cfg.use_phase_cmd: - self.phase = self._env.episode_length_buf[:, None] * self._env.step_dt / self.cycle_time - current_dist = torch.norm(self.robot.data.root_pos_w[:, :2] - self.env_origins[:, :2], dim=1) + current_dist = torch.norm(self.robot.data.root_pos_w[:, :2] - self._env.scene.env_origins[:, :2], dim=1) self.max_move_distance = torch.max(self.max_move_distance, current_dist) + def _update_env_command_ranges(self): + """ Update environment-wise command ranges based on current command ranges and terrain type """ + for terrain_type, terrain_command_ranges in self.cfg.terrain_max_command_ranges.items(): + if terrain_type not in self.terrain_type2idx: + continue + terrain_idx = self.terrain_type2idx[terrain_type] + env_ids = (self.terrain_idxs == terrain_idx).nonzero().flatten() + self.env_command_ranges['lin_vel_x'][env_ids, 0] = max( + terrain_command_ranges['lin_vel_x'][0], + self.command_ranges['lin_vel_x'][0], + ) + self.env_command_ranges['lin_vel_x'][env_ids, 1] = min( + terrain_command_ranges['lin_vel_x'][1], + self.command_ranges['lin_vel_x'][1] + ) + self.env_command_ranges['lin_vel_y'][env_ids, 0] = max( + terrain_command_ranges['lin_vel_y'][0], + self.command_ranges['lin_vel_y'][0] + ) + self.env_command_ranges['lin_vel_y'][env_ids, 1] = min( + terrain_command_ranges['lin_vel_y'][1], + self.command_ranges['lin_vel_y'][1] + ) + self.env_command_ranges['ang_vel_yaw'][env_ids, 0] = max( + terrain_command_ranges['ang_vel_yaw'][0], + self.command_ranges['ang_vel_yaw'][0] + ) + self.env_command_ranges['ang_vel_yaw'][env_ids, 1] = min( + terrain_command_ranges['ang_vel_yaw'][1], + self.command_ranges['ang_vel_yaw'][1] + ) + + def get_current_scale(self, config: dict): + """config: {'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0}""" + current_iter = self._env.common_step_counter // self.cfg.num_steps_per_iter + cfg_start_iter = config['start_iter'] + cfg_end_iter = config['end_iter'] + cfg_start_val = config['start_value'] + cfg_end_val = config['end_value'] + + percentage = (current_iter - cfg_start_iter) / (cfg_end_iter - cfg_start_iter) + percentage = max(min(percentage, 1.0), 0.0) + + current_scale = (1.0 - percentage) * cfg_start_val + percentage * cfg_end_val + return current_scale + + """Debug Visualization""" + def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility if debug_vis: - # create markers if necessary for the first time if not hasattr(self, "goal_vel_visualizer"): - # -- goal self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - # -- current self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - # set their visibility to true self.goal_vel_visualizer.set_visibility(True) self.current_vel_visualizer.set_visibility(True) else: @@ -346,52 +320,103 @@ class UniformVelTerrainCmd(CommandTerm): self.current_vel_visualizer.set_visibility(False) def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data if not self.robot.is_initialized: return - # get marker location - # -- base state base_pos_w = self.robot.data.root_pos_w.clone() base_pos_w[:, 2] += 0.5 - # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2]) - # display markers self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - """ - Internal helpers. - """ - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - # arrow-scale arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame base_quat_w = self.robot.data.root_quat_w arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat + def _resample_command(self): + ... @configclass -class UniformVelTerrainCmdCfg(UniformVelocityCommandCfg): +class Go2RLGymCommandCfg(CommandTermCfg): + class_type: type = Go2RLGymCommand - class_type: type = UniformVelTerrainCmd - cycle_time:float = 0.5 - max_angular_envs: float = 0.2 - limit_vel_envs: float = 0.2 - ranges: dict[str, UniformVelocityCommandCfg.Ranges] | UniformVelocityCommandCfg.Ranges = MISSING # 当前指令范围 - terrain_max_ranges: dict[str, UniformVelocityCommandCfg.Ranges] = MISSING # type:ignore 地形最大指令范围 - curriculum_schedule: list[dict] | None = None - use_phase_cmd: bool = True - \ No newline at end of file + asset_name: str = "robot" + """Name of the asset in the environment for which the commands are generated.""" + + dynamic_resample_commands: bool = True + """Sample commands with low bounds""" + limit_vel_invert_when_continuous: bool = True + """Invert the limit logic when using continuous sample limit velocity commands""" + + zero_command_curriculum: dict = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1} + """Start training with zero commands and then gradually increase zero command probability""" + limit_vel: dict = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} + """Sample vel commands from min [-1] or zero [0] or max [1] range only""" + command_range_curriculum: list[dict] = [{ + 'iter': 20000, # training iteration at which the command ranges are updated + 'lin_vel_x': [-1.0, 1.0], # min max [m/s] + 'lin_vel_y': [-1.0, 1.0], # min max [m/s] + 'ang_vel_yaw': [-1.5, 1.5], # min max [rad/s] + }, { + 'iter': 50000, # training iteration at which the command ranges are updated + 'lin_vel_x': [-2.0, 2.0], # min max [m/s] + 'lin_vel_y': [-1.0, 1.0], # min max [m/s] + 'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s] + }] + """List for command range curriculums at specific training iterations""" + terrain_max_command_ranges: dict[str, dict] = { + 'random_rough': + {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'hf_pyramid_slope': + {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'hf_pyramid_slope_inv': + {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'pyramid_stairs': + {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'pyramid_stairs_inv': + {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'boxes': + {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, + 'flat': + {'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0]}, + } + + resampling_time: float = 5.0 + resampling_time_range: tuple[float, float] = (5.0, 5.0) + """Time before command are changed [s]""" + limit_ang_vel_at_zero_command_prob: float = 0.2 + """Probability of add limiting angular velocity commands when zero command is sampled""" + limit_vel_prob: float = 0.2 + """Probability of limiting linear velocity command""" + num_steps_per_iter: int = 24 + """Number of envs steps for each training iteration""" + + @configclass + class Ranges: + lin_vel_x: tuple[float, float] = [-0.5, 0.5] + """Range for the linear-x velocity command [m/s]""" + lin_vel_y: tuple[float, float] = [-0.5, 0.5] + """Range for the linear-y velocity command [m/s]""" + ang_vel_yaw: tuple[float, float] = [-1.0, 1.0] + """Range for the angular-z velocity command [rad/s]""" + + ranges: Ranges = Ranges() + + goal_vel_visualizer_cfg: VisualizationMarkersCfg = GREEN_ARROW_X_MARKER_CFG.replace( + prim_path="/Visuals/Command/velocity_goal" + ) + """The configuration for the goal velocity visualization marker. Defaults to GREEN_ARROW_X_MARKER_CFG.""" + + current_vel_visualizer_cfg: VisualizationMarkersCfg = BLUE_ARROW_X_MARKER_CFG.replace( + prim_path="/Visuals/Command/velocity_current" + ) + """The configuration for the current velocity visualization marker. Defaults to BLUE_ARROW_X_MARKER_CFG.""" diff --git a/source/robot_lab/robot_lab/tasks/go2/mdp/commands_go2_rl_gym.py b/source/robot_lab/robot_lab/tasks/go2/mdp/commands_go2_rl_gym.py deleted file mode 100644 index 5f52e92..0000000 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/commands_go2_rl_gym.py +++ /dev/null @@ -1,422 +0,0 @@ -# -*- coding: utf-8 -*- -''' -@File : commands_go2_rl_gym.py -@Time : 2026/04/01 17:16:44 -@Author : wty-yy -@Version : 1.0 -@Blog : https://wty-yy.github.io/ -@Desc : CommandTerm for go2_rl_gym style command generation, reference to https://github.com/wty-yy/go2_rl_gym - -IsaacLab CommandTerm working flow: -Env: after compute reward call command.compute(dt) -1. self._update_metrics(): update self.metrics dict for logging -2. self.time_left -= dt -3. self._resample(self.time_left <= 0) -4. self._update_command(): update command if needed - -Get command from self.command property, return command, shape=(num_envs, command_dim) - -Note: -1. We don't use original self._resample(env_ids) and self._resample_command(env_ids), because it will randomize time_left -2. Remove heading command -3. We don't use curriculum item to update curriculum, inplace update -''' - -from __future__ import annotations # For forward reference of type hints - -from typing import TYPE_CHECKING, Sequence -if TYPE_CHECKING: # Avoid circular import for type checking - from robot_lab.tasks.go2.env.go2_env import ActionDelayGo2Env - -from itertools import product - -import torch - -from isaaclab.utils import configclass -from isaaclab.managers import CommandTerm, CommandTermCfg -from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg -from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG -from isaaclab.assets import Articulation -import isaaclab.utils.math as math_utils - -from robot_lab.tasks.go2.mdp.utils import is_robot_on_terrain, sample_disjoint_intervals, sample_single_interval - - -class Go2RLGymCommand(CommandTerm): - cfg: Go2RLGymCommandCfg - _env: ActionDelayGo2Env - - def __init__(self, cfg: Go2RLGymCommandCfg, env: ActionDelayGo2Env): - """Reference: https://github.com/wty-yy/go2_rl_gym/blob/master/legged_gym/envs/base/legged_robot.py - LeggedRobot._resample_command() and LeggedRobot._post_physics_step_callback() - """ - super().__init__(cfg, env) - self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, dtype=torch.float, device=self.device) - self.max_move_distance = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) - self.last_is_limit_vel = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - self.commands = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device) # [lin_vel_x, lin_vel_y, ang_vel_yaw] - self.command_ranges = self.cfg.ranges.to_dict() - self.env_command_ranges = { - 'lin_vel_x': torch.tensor(self.command_ranges['lin_vel_x'], device=self.device).repeat(self.num_envs, 1), - 'lin_vel_y': torch.tensor(self.command_ranges['lin_vel_y'], device=self.device).repeat(self.num_envs, 1), - 'ang_vel_yaw': torch.tensor(self.command_ranges['ang_vel_yaw'], device=self.device).repeat(self.num_envs, 1), - } - self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), - abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) - self.limit_vel_comb = torch.tensor(list(product( - self.cfg.limit_vel["lin_vel_x"], - self.cfg.limit_vel["lin_vel_y"], - self.cfg.limit_vel["ang_vel_yaw"] - )), device=self.device) - self._init_terrain_infos() - self._update_env_command_ranges() - self.robot: Articulation = env.scene[cfg.asset_name] - self.zero_command_prob = 0 - self.max_command_x = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) - - self.cfg.command_range_curriculum = sorted(self.cfg.command_range_curriculum, key=lambda x: x['iter'], reverse=True) - - def __str__(self) -> str: - """Return a string representation of the command term.""" - msg = (f"""Go2RLGymCommand:\n""" - f"""Command shape: {self.commands.shape}""") - return msg - - def _init_terrain_infos(self): - """Initialize terrain types and indices for each environment.""" - self.terrain_types = list(self._env.scene.terrain.cfg.terrain_generator.sub_terrains.keys()) - for terrain_type in self.terrain_types: - if terrain_type not in self.cfg.terrain_max_command_ranges: - raise ValueError(f"Terrain type '{terrain_type}' is not defined in cfg.terrain_max_command_ranges.") - self.terrain_type2idx = {terrain_type: idx for idx, terrain_type in enumerate(self.terrain_types)} - self.terrain_idxs = torch.full((self.num_envs,), -1, dtype=torch.long, device=self.device) - for terrain_type in self.terrain_types: - idxs = is_robot_on_terrain(self._env, terrain_type).nonzero().flatten() - if len(idxs) > 0: - self.terrain_idxs[idxs] = self.terrain_type2idx[terrain_type] - self.terrain_length = self._env.scene.terrain.cfg.terrain_generator.size[0] - - @property - def command(self) -> torch.Tensor: - return self.commands - - def _update_metrics(self): - self.max_command_x[:] = self.command_ranges["lin_vel_x"][1] - self.metrics["max_command_x"] = self.max_command_x - - def reset(self, env_ids: Sequence[int] | None = None): - self.time_left[env_ids] = self.cfg.resampling_time - self.commands_xy_accumulation[env_ids] = 0.0 - self.max_move_distance[env_ids] = 0.0 - return super().reset(env_ids) - - def _resample(self, env_ids: Sequence[int]): - """ Randommly select commands of some environments - - Args: - env_ids (List[int]): Environments ids for which new commands are needed - """ - env = self._env - if len(env_ids) == 0: - return - # update command curriculum with train steps - if len(self.cfg.command_range_curriculum): - current_iter = env.common_step_counter // self.cfg.num_steps_per_iter - for i in range(len(self.cfg.command_range_curriculum)-1, -1, -1): # iterate backwards to be able to pop entries - cfg = self.cfg.command_range_curriculum[i] - if current_iter >= cfg["iter"]: - self.command_ranges["lin_vel_x"] = cfg["lin_vel_x"] - self.command_ranges["lin_vel_y"] = cfg["lin_vel_y"] - self.command_ranges["ang_vel_yaw"] = cfg["ang_vel_yaw"] - self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), - abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) - self.cfg.command_range_curriculum.pop(i) - self._update_env_command_ranges() - print(f"Command range updated at iter {current_iter}: {self.command_ranges}") - remaining_dist = torch.clip(0.625 * self.terrain_length - torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.resampling_time, 0.0) - self.time_left[env_ids] = self.cfg.resampling_time - if self.cfg.dynamic_resample_commands: - # arrive at boundary 0.625 times the width of the remaining distance - if ((env.max_episode_length - env.episode_length_buf[env_ids]) + 1 == 0).any(): - raise ValueError("Some envs have zero remaining episode length during command resampling") - vel_low_bound = torch.clip(remaining_dist / ((env.max_episode_length - env.episode_length_buf[env_ids] + 1 + 1e-9) * env.step_dt), 0.0) - self.commands[env_ids, 0] = sample_disjoint_intervals( - env_ids, - vel_low_bound, - self.env_command_ranges["lin_vel_x"][env_ids, 0], - self.env_command_ranges["lin_vel_x"][env_ids, 1], - self.device - ) - self.commands[env_ids, 1] = sample_disjoint_intervals( - env_ids, - vel_low_bound, - self.env_command_ranges["lin_vel_y"][env_ids, 0], - self.env_command_ranges["lin_vel_y"][env_ids, 1], - self.device - ) - r = torch.rand(len(env_ids), device=self.device) - lower = self.env_command_ranges["ang_vel_yaw"][env_ids, 0] - upper = self.env_command_ranges["ang_vel_yaw"][env_ids, 1] - self.commands[env_ids, 2] = (upper - lower) * r + lower - else: - self.commands[env_ids, 0] = sample_single_interval( - env_ids, - self.env_command_ranges["lin_vel_x"][env_ids, 0], - self.env_command_ranges["lin_vel_x"][env_ids, 1], - self.device - ) - self.commands[env_ids, 1] = sample_single_interval( - env_ids, - self.env_command_ranges["lin_vel_y"][env_ids, 0], - self.env_command_ranges["lin_vel_y"][env_ids, 1], - self.device - ) - self.commands[env_ids, 2] = sample_single_interval( - env_ids, - self.env_command_ranges["ang_vel_yaw"][env_ids, 0], - self.env_command_ranges["ang_vel_yaw"][env_ids, 1], - self.device - ) - - # set small commands to zero - self.commands[env_ids, :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) - - rand_prob = torch.rand(len(env_ids), device=self.device) - min_prob, max_prob = 0.0, 0.0 - # set limitation lin vel - if self.cfg.limit_vel_prob > 0.0: - max_prob += self.cfg.limit_vel_prob - lim_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) - lim_env_ids = env_ids[lim_mask] - if len(lim_env_ids) > 0: - change_lim_env_ids = lim_env_ids - if self.cfg.limit_vel_invert_when_continuous: - was_limited = self.last_is_limit_vel[lim_env_ids] - invert_env_ids = lim_env_ids[was_limited] - self.commands[invert_env_ids, 0] *= -1.0 - self.commands[invert_env_ids, 1] *= -1.0 - self.commands[invert_env_ids, 2] *= -1.0 - change_lim_env_ids = lim_env_ids[~was_limited] - vel_idx = torch.randint(0, self.limit_vel_comb.shape[0], (len(change_lim_env_ids),), device=self.device) - lin_vel_x_lim = torch.where( - self.limit_vel_comb[vel_idx, 0] == -1, - self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 0], - self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 1], - ) - lin_vel_x_lim[self.limit_vel_comb[vel_idx, 0] == 0] = 0.0 - lin_vel_y_lim = torch.where( - self.limit_vel_comb[vel_idx, 1] == -1, - self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 0], - self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 1] - ) - lin_vel_y_lim[self.limit_vel_comb[vel_idx, 1] == 0] = 0.0 - ang_vel_z_lim = torch.where( - self.limit_vel_comb[vel_idx, 2] == -1, - self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 0], - self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 1] - ) - ang_vel_z_lim[self.limit_vel_comb[vel_idx, 2] == 0] = 0.0 - self.commands[change_lim_env_ids, 0] = lin_vel_x_lim - self.commands[change_lim_env_ids, 1] = lin_vel_y_lim - self.commands[change_lim_env_ids, 2] = ang_vel_z_lim - self.last_is_limit_vel[env_ids] = False - self.last_is_limit_vel[lim_env_ids] = True - else: - self.last_is_limit_vel[env_ids] = False - min_prob += self.cfg.limit_vel_prob - - # set all commands to zero with some probability - if self.cfg.zero_command_curriculum is not None: - self.zero_command_prob = self.get_current_scale(self.cfg.zero_command_curriculum) - if self.zero_command_prob > 0.0: - max_prob += self.zero_command_prob - next_time_left = torch.clip( - env.max_episode_length_s - env.episode_length_buf[env_ids] * env.step_dt - (remaining_dist / (0.8 * self.max_lin_vel + 1e-9)), - min=0.0, - max=self.cfg.resampling_time, - ) - zero_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) * (next_time_left > 0.0) - zero_env_ids = env_ids[zero_mask] - if len(zero_env_ids) > 0: - self.commands[zero_env_ids, :2] = 0.0 - self.time_left[zero_env_ids] = next_time_left[zero_mask] - if self.cfg.limit_ang_vel_at_zero_command_prob > 0.0: - ang_vel_rand = torch.rand(len(zero_env_ids), device=self.device) # independent distribution - add_ang_mask = ang_vel_rand < self.cfg.limit_ang_vel_at_zero_command_prob - add_ang_env_ids = zero_env_ids[add_ang_mask] - if len(add_ang_env_ids) > 0: - direction_rand = torch.rand(len(add_ang_env_ids), device=self.device) - self.commands[add_ang_env_ids, 2] = torch.where( - direction_rand < 0.5, - self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 0], - self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 1] - ) - min_prob += self.zero_command_prob - - self.commands_xy_accumulation[env_ids] += self.commands[env_ids, :2] - - def _update_command(self): - current_dist = torch.norm(self.robot.data.root_pos_w[:, :2] - self._env.scene.env_origins[:, :2], dim=1) - self.max_move_distance = torch.max(self.max_move_distance, current_dist) - - def _update_env_command_ranges(self): - """ Update environment-wise command ranges based on current command ranges and terrain type """ - for terrain_type, terrain_command_ranges in self.cfg.terrain_max_command_ranges.items(): - if terrain_type not in self.terrain_type2idx: - continue - terrain_idx = self.terrain_type2idx[terrain_type] - env_ids = (self.terrain_idxs == terrain_idx).nonzero().flatten() - self.env_command_ranges['lin_vel_x'][env_ids, 0] = max( - terrain_command_ranges['lin_vel_x'][0], - self.command_ranges['lin_vel_x'][0], - ) - self.env_command_ranges['lin_vel_x'][env_ids, 1] = min( - terrain_command_ranges['lin_vel_x'][1], - self.command_ranges['lin_vel_x'][1] - ) - self.env_command_ranges['lin_vel_y'][env_ids, 0] = max( - terrain_command_ranges['lin_vel_y'][0], - self.command_ranges['lin_vel_y'][0] - ) - self.env_command_ranges['lin_vel_y'][env_ids, 1] = min( - terrain_command_ranges['lin_vel_y'][1], - self.command_ranges['lin_vel_y'][1] - ) - self.env_command_ranges['ang_vel_yaw'][env_ids, 0] = max( - terrain_command_ranges['ang_vel_yaw'][0], - self.command_ranges['ang_vel_yaw'][0] - ) - self.env_command_ranges['ang_vel_yaw'][env_ids, 1] = min( - terrain_command_ranges['ang_vel_yaw'][1], - self.command_ranges['ang_vel_yaw'][1] - ) - - def get_current_scale(self, config: dict): - """config: {'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0}""" - current_iter = self._env.common_step_counter // self.cfg.num_steps_per_iter - cfg_start_iter = config['start_iter'] - cfg_end_iter = config['end_iter'] - cfg_start_val = config['start_value'] - cfg_end_val = config['end_value'] - - percentage = (current_iter - cfg_start_iter) / (cfg_end_iter - cfg_start_iter) - percentage = max(min(percentage, 1.0), 0.0) - - current_scale = (1.0 - percentage) * cfg_start_val + percentage * cfg_end_val - return current_scale - - """Debug Visualization""" - - def _set_debug_vis_impl(self, debug_vis: bool): - if debug_vis: - if not hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - self.goal_vel_visualizer.set_visibility(True) - self.current_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - if not self.robot.is_initialized: - return - base_pos_w = self.robot.data.root_pos_w.clone() - base_pos_w[:, 2] += 0.5 - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2]) - self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) - self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - base_quat_w = self.robot.data.root_quat_w - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - - def _resample_command(self): - ... - -@configclass -class Go2RLGymCommandCfg(CommandTermCfg): - class_type: type = Go2RLGymCommand - - asset_name: str = "robot" - """Name of the asset in the environment for which the commands are generated.""" - - dynamic_resample_commands: bool = True - """Sample commands with low bounds""" - limit_vel_invert_when_continuous: bool = True - """Invert the limit logic when using continuous sample limit velocity commands""" - - zero_command_curriculum: dict = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1} - """Start training with zero commands and then gradually increase zero command probability""" - limit_vel: dict = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} - """Sample vel commands from min [-1] or zero [0] or max [1] range only""" - command_range_curriculum: list[dict] = [{ - 'iter': 20000, # training iteration at which the command ranges are updated - 'lin_vel_x': [-1.0, 1.0], # min max [m/s] - 'lin_vel_y': [-1.0, 1.0], # min max [m/s] - 'ang_vel_yaw': [-1.5, 1.5], # min max [rad/s] - }, { - 'iter': 50000, # training iteration at which the command ranges are updated - 'lin_vel_x': [-2.0, 2.0], # min max [m/s] - 'lin_vel_y': [-1.0, 1.0], # min max [m/s] - 'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s] - }] - """List for command range curriculums at specific training iterations""" - terrain_max_command_ranges: dict[str, dict] = { - 'random_rough': - {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'hf_pyramid_slope': - {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'hf_pyramid_slope_inv': - {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'pyramid_stairs': - {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'pyramid_stairs_inv': - {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'boxes': - {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5]}, - 'flat': - {'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0]}, - } - - resampling_time: float = 5.0 - resampling_time_range: tuple[float, float] = (5.0, 5.0) - """Time before command are changed [s]""" - limit_ang_vel_at_zero_command_prob: float = 0.2 - """Probability of add limiting angular velocity commands when zero command is sampled""" - limit_vel_prob: float = 0.2 - """Probability of limiting linear velocity command""" - num_steps_per_iter: int = 24 - """Number of envs steps for each training iteration""" - - @configclass - class Ranges: - lin_vel_x: tuple[float, float] = [-0.5, 0.5] - """Range for the linear-x velocity command [m/s]""" - lin_vel_y: tuple[float, float] = [-0.5, 0.5] - """Range for the linear-y velocity command [m/s]""" - ang_vel_yaw: tuple[float, float] = [-1.0, 1.0] - """Range for the angular-z velocity command [rad/s]""" - - ranges: Ranges = Ranges() - - goal_vel_visualizer_cfg: VisualizationMarkersCfg = GREEN_ARROW_X_MARKER_CFG.replace( - prim_path="/Visuals/Command/velocity_goal" - ) - """The configuration for the goal velocity visualization marker. Defaults to GREEN_ARROW_X_MARKER_CFG.""" - - current_vel_visualizer_cfg: VisualizationMarkersCfg = BLUE_ARROW_X_MARKER_CFG.replace( - prim_path="/Visuals/Command/velocity_current" - ) - """The configuration for the current velocity visualization marker. Defaults to BLUE_ARROW_X_MARKER_CFG."""