📝 Add docstrings to g1-energy-benchmark-3027240647789121151

Docstrings generation was requested by @ngoiyaeric.

* #1 (comment)

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`

Author: coderabbitai[bot]

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

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")

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)

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),
+    )

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