From cacdffb622728d1fe05ef9564c4b3a481fdebf4c Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Sun, 1 Mar 2026 17:13:16 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20Add=20docstrings=20to=20`g1-ener?= =?UTF-8?q?gy-benchmark-3027240647789121151`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Docstrings generation was requested by @ngoiyaeric. * https://github.com/ngoiyaeric/IsaacLab/pull/1#issuecomment-3980291831 The following files were modified: * `source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py` * `source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py` * `source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py` * `source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.py` * `source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py` --- .../velocity/config/g1_energy/env.py | 189 ++++++++++++++++++ .../velocity/config/g1_energy/env_cfg.py | 62 ++++++ .../config/g1_energy/mdp/observations.py | 33 +++ .../velocity/config/g1_energy/mdp/rewards.py | 42 ++++ .../config/g1_energy/mdp/terminations.py | 21 ++ 5 files changed, 347 insertions(+) create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py new file mode 100644 index 00000000000..cc0245a6b90 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py @@ -0,0 +1,189 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.envs.manager_based_rl_env import ManagerBasedRLEnv + + +class G1EnergyEnv(ManagerBasedRLEnv): + """ + Custom environment for the G1 robot with an energy benchmark. + This intercepts the environment step to track battery state and tokens. + """ + + def __init__(self, cfg, **kwargs): + # 1. Allocate custom buffers before calling super().__init__() + # These will be initialized to 1.0 (full battery) and 0.0 (no tokens) + """ + Initialize the G1EnergyEnv and prepare placeholders for per-environment energy buffers. + + Initializes attributes used by the energy benchmark: `battery_buf` and `tokens_buf` are set to None here and will be allocated and initialized inside `load_managers` (buffers are initialized to full battery and zero tokens once the number of environments and device are known). Stores the provided configuration in `_cfg` and then delegates further initialization to the parent class by calling `super().__init__`, which triggers `load_managers`. + + Parameters: + cfg: Configuration object or mapping containing environment settings used by this class. Expected keys used later include `battery_capacity`, `battery_drain_rate`, `token_earn_rate`, `charge_token_cost`, and `charging_station_radius`. + **kwargs: Additional keyword arguments forwarded to the parent class initializer. + """ + self.battery_buf = None + self.tokens_buf = None + self._cfg = cfg + + # Super init will call load_managers, so we will initialize the buffers inside load_managers + super().__init__(cfg, **kwargs) + + def load_managers(self): + """ + Prepare per-environment energy state and load energy-related configuration. + + Initializes `battery_buf` to ones and `tokens_buf` to zeros with length `self.num_envs` on `self.device`, and stores energy parameters from `self.cfg`: `max_battery` (battery_capacity), `battery_drain_rate`, `token_earn_rate`, `charge_token_cost`, and `charging_station_radius`. + """ + super().load_managers() + # Initialize the buffers now that the number of environments is known + self.battery_buf = torch.ones(self.num_envs, device=self.device) + self.tokens_buf = torch.zeros(self.num_envs, device=self.device) + + # Parameters for energy / tokens + self.max_battery = self.cfg.battery_capacity + self.battery_drain_rate = self.cfg.battery_drain_rate + self.token_earn_rate = self.cfg.token_earn_rate + self.charge_token_cost = self.cfg.charge_token_cost + self.charging_station_radius = self.cfg.charging_station_radius + + def step(self, action: torch.Tensor): + # Process actions + """ + Advance the environment one control step using the provided actions, update energy/tokens state, run managers, handle resets, and return the latest observations and metrics. + + Parameters: + action (torch.Tensor): Actions for all environments; will be applied to the simulation. + + Returns: + tuple: A 5-tuple (obs_buf, reward_buf, reset_terminated, reset_time_outs, extras) where + - obs_buf: latest observation buffer for all environments, + - reward_buf: reward values computed for this step, + - reset_terminated: boolean mask indicating environments reset due to termination, + - reset_time_outs: boolean mask indicating environments reset due to timeout, + - extras: dictionary of additional info and logged metrics. + """ + self.action_manager.process_action(action.to(self.device)) + self.recorder_manager.record_pre_step() + + is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors() + + # Perform physics stepping + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + self.action_manager.apply_action() + self.scene.write_data_to_sim() + self.sim.step(render=False) + self.recorder_manager.record_post_physics_decimation_step() + + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render() + + self.scene.update(dt=self.physics_dt) + + # -- UPDATE BUFFERS AND ENERGY/TOKENS -- + self.episode_length_buf += 1 + self.common_step_counter += 1 + + # Calculate energy drain based on power consumption + robot = self.scene["robot"] + + # Using simplified energy drain: sum of absolute torques + # You could also use the actual power formula: |torque * velocity| + energy_drain = torch.sum(torch.abs(robot.data.applied_torque), dim=1) * self.battery_drain_rate * self.step_dt + self.battery_buf = torch.clamp(self.battery_buf - energy_drain, min=0.0, max=self.max_battery) + + # Calculate token earning (Job: Tracking velocity) + # Job quality depends on linear velocity tracking error + vel_cmd = self.command_manager.get_command("base_velocity") + current_vel = robot.data.root_lin_vel_b + + vel_error = torch.sum(torch.square(vel_cmd[:, :2] - current_vel[:, :2]), dim=1) + + # Earn tokens if error is low (robot is doing its job well) + job_quality = torch.exp(-vel_error / 0.5) + tokens_earned = job_quality * self.token_earn_rate * self.step_dt + self.tokens_buf += tokens_earned + + # Charging logic + # Check distance to charging station (origin 0,0) + dist_to_station = torch.norm(robot.data.root_pos_w[:, :2], dim=1) + at_station = dist_to_station < self.charging_station_radius + can_charge = self.tokens_buf >= self.charge_token_cost + + charging_envs = at_station & can_charge + + if charging_envs.any(): + charging_ids = charging_envs.nonzero(as_tuple=False).flatten() + + # Apply charge + self.battery_buf[charging_ids] = self.max_battery + self.tokens_buf[charging_ids] -= self.charge_token_cost + + # Trigger custom event for visual / logging if needed + if "at_charging_station" in self.event_manager.available_modes: + self.event_manager.apply(mode="at_charging_station", env_ids=charging_ids) + + # -- MANAGERS -- + self.reset_buf = self.termination_manager.compute() + self.reset_terminated = self.termination_manager.terminated + self.reset_time_outs = self.termination_manager.time_outs + + self.reward_buf = self.reward_manager.compute(dt=self.step_dt) + + if len(self.recorder_manager.active_terms) > 0: + self.obs_buf = self.observation_manager.compute() + self.recorder_manager.record_post_step() + + # Reset environments + reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + if len(reset_env_ids) > 0: + self.recorder_manager.record_pre_reset(reset_env_ids) + self._reset_idx(reset_env_ids) + + if self.sim.has_rtx_sensors() and self.cfg.num_rerenders_on_reset > 0: + for _ in range(self.cfg.num_rerenders_on_reset): + self.sim.render() + + self.recorder_manager.record_post_reset(reset_env_ids) + + # -- COMMANDS & EVENTS -- + self.command_manager.compute(dt=self.step_dt) + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) + + self.obs_buf = self.observation_manager.compute(update_history=True) + + return ( + self.obs_buf, + self.reward_buf, + self.reset_terminated, + self.reset_time_outs, + self.extras, + ) + + def _reset_idx(self, env_ids: torch.Tensor): + """ + Reset the specified environments' energy-related state and record average metrics. + + Resets battery level to full and tokens to zero for the given environment indices, then stores the average battery and token values in extras["log"] under "Metrics/avg_battery" and "Metrics/avg_tokens". + + Parameters: + env_ids (torch.Tensor): 1D tensor of environment indices to reset. + """ + super()._reset_idx(env_ids) + + # Reset custom buffers + self.battery_buf[env_ids] = self.max_battery + self.tokens_buf[env_ids] = 0.0 + + # Add custom metrics logging + avg_battery = torch.mean(self.battery_buf).item() + avg_tokens = torch.mean(self.tokens_buf).item() + + self.extras["log"]["Metrics/avg_battery"] = avg_battery + self.extras["log"]["Metrics/avg_tokens"] = avg_tokens diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py new file mode 100644 index 00000000000..d201f618599 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py @@ -0,0 +1,62 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.locomotion.velocity.config.g1.flat_env_cfg import ( + G1FlatEnvCfg, +) + +import source.isaaclab_tasks.isaaclab_tasks.manager_based.locomotion.velocity.config.g1_energy.mdp as custom_mdp + + +@configclass +class G1EnergyEnvCfg(G1FlatEnvCfg): + """Configuration for the G1 Energy Environment.""" + + # Energy / Token configurations + battery_capacity: float = 1.0 + battery_drain_rate: float = 0.005 # Rate of battery drain based on torque + token_earn_rate: float = 0.1 # Tokens earned per second of good tracking + charge_token_cost: float = 1.0 # Cost in tokens to fully charge + charging_station_radius: float = 1.0 # Radius to trigger charging at origin + + def __post_init__(self): + """ + Finalize post-initialization for the energy environment configuration by applying energy-specific command limits, episode length, terminations, rewards, observations, and event tracking. + + This method adjusts the base movement command ranges, extends the episode duration to accommodate charging behavior, registers a battery-empty termination, adds battery-related reward terms (including an empty-battery penalty), exposes battery level and token count as policy observations, and enables tracking of the "at_charging_station" event. + """ + super().__post_init__() + + # Overwrite the base environment commands + # Allow the robot to track X, Y and Yaw + self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) + self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5) + self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) + + # Extend episode length to allow for longer charging cycles + self.episode_length_s = 60.0 + + # Terminations + self.terminations.battery_empty = DoneTerm(func=custom_mdp.battery_empty, time_out=True) + + # Rewards + self.rewards.battery_penalty = RewTerm( + func=custom_mdp.battery_penalty, weight=0.1 + ) # Note: returns negative inside + self.rewards.empty_battery_penalty = RewTerm(func=custom_mdp.empty_battery_penalty, weight=1.0) + + # Observations + # Add the custom observation terms to the policy observation space + self.observations.policy.battery_level = ObsTerm(func=custom_mdp.battery_level) + self.observations.policy.token_count = ObsTerm(func=custom_mdp.token_count) + + # Ensure event mode is tracked + self.events.at_charging_station = EventTerm(func=lambda env, env_ids: None, mode="at_charging_station") diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py new file mode 100644 index 00000000000..b9cf6a01177 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.envs.manager_based_rl_env import ManagerBasedRLEnv + + +def battery_level(env: ManagerBasedRLEnv) -> torch.Tensor: + """ + Get the robot's battery level normalized to its maximum capacity. + + Parameters: + env (ManagerBasedRLEnv): Environment containing `battery_buf` and `max_battery`. + + Returns: + torch.Tensor: Tensor of shape (num_envs, 1) with values in [0, 1] representing each environment's battery level divided by `max_battery`. + """ + # (num_envs, 1) + return (env.battery_buf / env.max_battery).unsqueeze(-1) + + +def token_count(env: ManagerBasedRLEnv) -> torch.Tensor: + """ + Provide the current token count per environment as a single-column tensor. + + Returns: + torch.Tensor: Tensor of shape (num_envs, 1) containing the token count for each environment. + """ + # (num_envs, 1) + return env.tokens_buf.unsqueeze(-1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.py new file mode 100644 index 00000000000..56781fa072f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.envs.manager_based_rl_env import ManagerBasedRLEnv + + +def battery_penalty(env: ManagerBasedRLEnv) -> torch.Tensor: + """ + Compute a linear penalty proportional to remaining battery. + + Parameters: + env (ManagerBasedRLEnv): Environment providing `battery_buf` and `max_battery`. + + Returns: + torch.Tensor: Penalty computed as -(1.0 - (env.battery_buf / env.max_battery)); values are 0 when battery is full and approach -1 as battery approaches empty. + """ + # Scale from 0 (full) to -1 (empty) linearly + # You could make it non-linear e.g., only penalize if below 20% + return -(1.0 - (env.battery_buf / env.max_battery)) + + +def empty_battery_penalty(env: ManagerBasedRLEnv) -> torch.Tensor: + """ + Apply a heavy penalty when the environment's battery is effectively empty. + + Parameters: + env (ManagerBasedRLEnv): Environment exposing `battery_buf` (current battery level), + `device` (tensor device), and `max_battery` (not used here). + + Returns: + torch.Tensor: A scalar tensor on `env.device` with value `-10.0` if `env.battery_buf <= 0.01`, + `0.0` otherwise. + """ + return torch.where( + env.battery_buf <= 0.01, + torch.tensor(-10.0, device=env.device), + torch.tensor(0.0, device=env.device), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py new file mode 100644 index 00000000000..40433488235 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.envs.manager_based_rl_env import ManagerBasedRLEnv + + +def battery_empty(env: ManagerBasedRLEnv) -> torch.Tensor: + """ + Indicates whether an episode should terminate when the environment's battery level is depleted. + + Parameters: + env (ManagerBasedRLEnv): Environment whose `battery_buf` tensor is checked for depletion. + + Returns: + torch.Tensor: Boolean tensor with `True` where `battery_buf` is less than or equal to 0.0, `False` otherwise. + """ + return env.battery_buf <= 0.0