feat: Introduce G1 Energy Benchmark Environment (Fluidity Index)

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/__init__.py

@@ -5,5 +5,5 @@
 
 """Configurations for velocity-based locomotion environments."""
 
-# We leave this file empty since we don't want to expose any configs in this package directly.
-# We still need this file to import the "config" module in the parent package.
+# Re-export g1_energy configs for parent-package imports
+from .g1_energy import *

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/__init__.py

@@ -0,0 +1,19 @@
+# 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 gymnasium as gym
+
+from .env import G1EnergyEnv
+from .env_cfg import G1EnergyEnvCfg
+
+gym.register(
+    id="Isaac-Velocity-Energy-G1-v0",
+    entry_point="isaaclab_tasks.manager_based.locomotion.velocity.config.g1_energy:G1EnergyEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.env_cfg:G1EnergyEnvCfg",
+        "rsl_rl_cfg_entry_point": "isaaclab_tasks.manager_based.locomotion.velocity.config.g1.agents.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg",
+    },
+)

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env.py

@@ -0,0 +1,156 @@
+# 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)
+        self.battery_buf = None
+        self.tokens_buf = None
+
+        # Super init will call load_managers, so we will initialize the buffers inside load_managers
+        super().__init__(cfg, **kwargs)
+
+    def load_managers(self):
+        # 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
+
+        # Call super now that max_battery, battery_buf, and tokens_buf are initialized
+        # so that when ObservationManager runs during initialization, it has access to them
+        super().load_managers()
+
+    def step(self, action: torch.Tensor):
+        # Process actions
+        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 / self.cfg.vel_error_scale)
+        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 locally)
+        pos_local_xy = robot.data.root_pos_w[:, :2] - self.scene.env_origins[:, :2]
+        dist_to_station = torch.norm(pos_local_xy, 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).flatten()
+        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):
+        # Sample custom metrics logging before buffer reset
+        avg_battery = torch.mean(self.battery_buf[env_ids]).item()
+        avg_tokens = torch.mean(self.tokens_buf[env_ids]).item()
+
+        super()._reset_idx(env_ids)
+
+        self.extras["log"]["Metrics/avg_battery"] = avg_battery
+        self.extras["log"]["Metrics/avg_tokens"] = avg_tokens
+
+        # Reset custom buffers
+        self.battery_buf[env_ids] = self.max_battery
+        self.tokens_buf[env_ids] = 0.0

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/env_cfg.py

@@ -0,0 +1,63 @@
+# 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
+
+import isaaclab_tasks.manager_based.locomotion.velocity.config.g1_energy.mdp as custom_mdp
+from isaaclab_tasks.manager_based.locomotion.velocity.config.g1.flat_env_cfg import (
+    G1FlatEnvCfg,
+)
+
+
+def _noop_event(env, env_ids):
+    pass
+
+
+@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
+    vel_error_scale: float = 0.5  # Scale applied to tracking error for earning tokens
+
+    def __post_init__(self):
+        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=False)
+
+        # 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, params={"max_expected_tokens": self.charge_token_cost}
+        )
+
+        # Ensure event mode is tracked
+        self.events.at_charging_station = EventTerm(func=_noop_event, mode="at_charging_station")

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/__init__.py

@@ -0,0 +1,8 @@
+# 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 .observations import *
+from .rewards import *
+from .terminations import *

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/observations.py

@@ -0,0 +1,20 @@
+# 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:
+    """Returns the current normalized battery level of the robot."""
+    # (num_envs, 1)
+    return (env.battery_buf / env.max_battery).unsqueeze(-1)
+
+
+def token_count(env: ManagerBasedRLEnv, max_expected_tokens: float = 1.0) -> torch.Tensor:
+    """Returns the current normalized token count of the robot."""
+    # (num_envs, 1)
+    return (env.tokens_buf / max_expected_tokens).unsqueeze(-1)

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/rewards.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_penalty(env: ManagerBasedRLEnv) -> torch.Tensor:
+    """Penalize the robot when its battery gets too low."""
+    # 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:
+    """Heavily penalize when battery hits exactly 0."""
+    mask = env.battery_buf <= 0.01
+    return mask.to(env.battery_buf.dtype) * -10.0

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/terminations.py

@@ -0,0 +1,13 @@
+# 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:
+    """Terminate the episode if the battery level reaches 0."""
+    return env.battery_buf <= 0.0