diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/__init__.py index 26f3257daef..1c25f751794 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/__init__.py +++ b/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 * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/__init__.py new file mode 100644 index 00000000000..a85c81dcb72 --- /dev/null +++ b/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", + }, +) 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..cd92c5f1787 --- /dev/null +++ b/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 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..e82e044fc48 --- /dev/null +++ b/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") diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1_energy/mdp/__init__.py new file mode 100644 index 00000000000..fd7240379ea --- /dev/null +++ b/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 * 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..071d059bebe --- /dev/null +++ b/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) 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..ed5197be968 --- /dev/null +++ b/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 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..e862310f280 --- /dev/null +++ b/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