"""Task-owned Manager-Based terms for G1 biped locomotion.
The reward/observation equations and tuned constants are ported from UniLab's
pre-Manager-Based G1Walk joystick task family (ported in 64c7f8b8, roadmap
#1042). The adaptation uses community ``func + params`` terms, Torch state
carriers, and cached cold-path sensor bindings; hot paths never parse assets
or probe backend privates.
"""
from __future__ import annotations
import math
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast
from weakref import WeakKeyDictionary
import torch
from unilab.base.backend_factory import create_backend, env_backend_kwargs
from unilab.base.curriculum import EpisodeLengthTracker
from unilab.base.reset_selectors import ResetMaskBuffer
from unilab.envs.manager_based_rl_env import (
ManagerBasedRlEnv as _ConcreteManagerBasedRlEnv,
)
from unilab.envs.manager_based_rl_env import (
ManagerBasedRlEnvCfg,
_resolve_backend_entity_contract,
)
from unilab.envs.mdp.commands.velocity_command import (
UniformVelocityCommand,
UniformVelocityCommandCfg,
)
from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg
from unilab.managers.reset_owner import ResetOwner, ResetOwnerCfg
from unilab.managers.reward_manager import RewardTermCfg
from unilab.managers.rng import ManagerRng
from unilab.managers.scene_entity_config import SceneEntityCfg
from unilab.managers.termination_manager import TerminationTermCfg
from unilab.tasks.locomotion.common.manager_terms import SensorTermBase
if TYPE_CHECKING:
from unilab.base.entity import Entity
from unilab.managers.types import ManagerBasedRlEnv
class _RewardTermCfgView(Protocol):
weight: float
class _RewardManagerView(Protocol):
@property
def active_terms(self) -> list[str]: ...
def get_term_cfg(self, term_name: str) -> _RewardTermCfgView: ...
class _G1Env(ManagerBasedRlEnv, Protocol):
@property
def common_step_counter(self) -> int: ...
@property
def reset_buf(self) -> torch.Tensor: ...
@property
def reward_manager(self) -> _RewardManagerView: ...
@property
def cfg(self) -> Any: ...
@property
def scene(self) -> Any: ...
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
_LEFT_FOOT_CONTACT_SENSORS = tuple(f"left_foot_contact_{index}" for index in range(4))
_RIGHT_FOOT_CONTACT_SENSORS = tuple(f"right_foot_contact_{index}" for index in range(4))
_FOOT_CONTACT_SENSORS = _LEFT_FOOT_CONTACT_SENSORS + _RIGHT_FOOT_CONTACT_SENSORS
_FOOT_POS_SENSORS = ("left_foot_pos", "right_foot_pos")
_FOOT_QUAT_SENSORS = ("left_foot_quat", "right_foot_quat")
_GAIT_INIT_MODES = ("offset_phase", "independent")
def _real(
term: str,
name: str,
value: Any,
*,
minimum: float | None = None,
strict_minimum: bool = False,
) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"{term} {name} must be a real number")
result = float(value)
if not math.isfinite(result):
raise ValueError(f"{term} {name} must be finite")
if minimum is not None and (result <= minimum if strict_minimum else result < minimum):
relation = "greater than" if strict_minimum else "at least"
raise ValueError(f"{term} {name} must be {relation} {minimum}")
return result
def _weights(term: str, name: str, value: Any) -> torch.Tensor:
if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)):
raise TypeError(f"{term} {name} must be a sequence of real numbers")
result = torch.tensor(
[_real(term, f"{name}[{index}]", item) for index, item in enumerate(value)],
dtype=torch.float32,
)
if result.ndim != 1 or result.numel() == 0:
raise ValueError(f"{term} {name} must be a non-empty one-dimensional sequence")
return result
def _square(value: torch.Tensor) -> torch.Tensor:
return torch.square(_torch_values(value))
def _exp(value: torch.Tensor) -> torch.Tensor:
return torch.exp(_torch_values(value))
def _torch_values(value: torch.Tensor) -> torch.Tensor:
return value.to(dtype=torch.float32)
def _command(env: ManagerBasedRlEnv, term: str, command_name: str) -> torch.Tensor:
if not isinstance(command_name, str) or not command_name:
raise ValueError(f"{term} command_name must be a non-empty string")
try:
command = env.command_manager.get_command(command_name)
except KeyError as exc:
raise KeyError(f"{term} command capability '{command_name}' is unavailable") from exc
if command is None:
raise KeyError(f"{term} command capability '{command_name}' is unavailable")
command_tensor = _torch_values(command)
if tuple(command_tensor.shape) != (env.num_envs, 3):
raise ValueError(
f"{term} command '{command_name}' must have shape "
f"{(env.num_envs, 3)}; got {tuple(command_tensor.shape)}"
)
return command_tensor
def _asset(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg) -> Entity:
return cast("Entity", env.scene[asset_cfg.name])
class _SensorTerm(SensorTermBase):
"""G1 manager terms share the locomotion-wide named-sensor binding."""
def _read_state_tensor(
self,
env: ManagerBasedRlEnv,
capability: str,
binding: tuple[Any, tuple[str, ...]],
shape: tuple[int, ...],
) -> torch.Tensor:
view, sensor_names = binding
read_plan = getattr(env.scene, "read_plan", None)
owner_sensors = set(getattr(read_plan, "sensor_names", {}).get("robot", ()))
if read_plan is None or not set(sensor_names).issubset(owner_sensors):
raise NotImplementedError(
f"Manager term '{self.name}' requires its {capability} sensors "
f"{sensor_names} in the packed Torch read plan"
)
entity = env.scene["robot"]
views = read_plan.sensor_tensor_views(entity, sensor_names).values
if len(sensor_names) == 1:
values = views[sensor_names[0]]
if tuple(values.shape) != shape:
raise ValueError(
f"{self.name} {capability} must have shape {shape}, got {tuple(values.shape)}"
)
return values
values = torch.cat([views[name] for name in sensor_names], dim=1)
if tuple(values.shape) != shape:
raise ValueError(
f"{self.name} {capability} must have shape {shape}, got {tuple(values.shape)}"
)
return values
# ---------------------------------------------------------------------------
# Gait phase state (shared by the observation term and the gait reward terms)
# ---------------------------------------------------------------------------
[docs]
def compute_feet_phase_height_targets(
gait_phase: torch.Tensor, swing_height: float
) -> tuple[torch.Tensor, torch.Tensor]:
"""Cubic-Bézier per-foot height targets, from the G1Walk task family (64c7f8b8)."""
def cubic_bezier_height(phase: torch.Tensor, swing_height: float) -> torch.Tensor:
phase = phase if phase.ndim == 2 else phase[:, None]
phi_normalized = torch.remainder(phase + math.pi, 2 * math.pi) - math.pi
x = (phi_normalized + math.pi) / (2 * math.pi)
stance_t = (2 * x).clamp(max=1.0)
swing_t = (2 * x - 1).clamp(min=0.0)
bezier = torch.where(
x <= 0.5,
stance_t**3 + 3 * stance_t**2 * (1 - stance_t),
1.0 - (swing_t**3 + 3 * swing_t**2 * (1 - swing_t)),
)
return swing_height * bezier
targets = cubic_bezier_height(_torch_values(gait_phase), swing_height)
return targets[:, 0], targets[:, 1]
@dataclass
class _G1GaitContext:
phase: torch.Tensor # (num_envs, 2), radians in [0, 2*pi)
delta: float # 2*pi*frequency*ctrl_dt
frequency: float
init_mode: str
last_counter: int
_GAIT_CONTEXTS: WeakKeyDictionary[Any, _G1GaitContext] = WeakKeyDictionary()
_GAIT_PHASE_CACHE: WeakKeyDictionary[Any, tuple[_G1GaitContext, int]] = WeakKeyDictionary()
def _gait_context(env: _G1Env, term: str, frequency: float, init_mode: str) -> _G1GaitContext:
context = _GAIT_CONTEXTS.get(env)
if context is None:
delta = float(
2.0
* math.pi
* frequency
* _real(term, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True)
)
device = torch.device(env.device)
phase = torch.zeros((env.num_envs, 2), dtype=torch.float32, device=device)
context = _G1GaitContext(
phase=phase,
delta=delta,
frequency=frequency,
init_mode=init_mode,
last_counter=int(env.common_step_counter),
)
_GAIT_CONTEXTS[env] = context
return context
if context.frequency != frequency or context.init_mode != init_mode:
raise ValueError(
f"{term} gait terms on one env must share frequency and init_mode; "
f"got frequency {context.frequency} vs {frequency}, "
f"init_mode {context.init_mode!r} vs {init_mode!r}"
)
return context
def _advance_gait(env: _G1Env, context: _G1GaitContext) -> torch.Tensor:
counter = env.common_step_counter
if isinstance(counter, bool) or not isinstance(counter, int):
raise TypeError("G1 gait terms require an integer common_step_counter")
counter = int(counter)
if counter < context.last_counter:
raise ValueError("G1 gait terms common_step_counter cannot move backwards")
steps = counter - context.last_counter
if steps:
context.phase = torch.remainder(context.phase + steps * context.delta, 2.0 * math.pi)
context.last_counter = counter
return context.phase
def _cached_gait_phase(env: _G1Env, context: _G1GaitContext) -> torch.Tensor:
"""Return the phase already advanced for this env and step.
Gait rewards are semantically pure in the phase for a given counter:
``_advance_gait`` is counter-idempotent, but its earlier host-side dispatch
and tensor ops were repeated by every gait reward. The cache is keyed by
the same weak env identity and exact context object, and reset only changes
selected rows after the latest public phase read. A backward counter is
still rejected by ``_advance_gait`` on the first call of a new step.
"""
cached = _GAIT_PHASE_CACHE.get(env)
counter = int(env.common_step_counter)
if cached is not None:
cached_context, cached_counter = cached
if cached_context is context and cached_counter == counter:
return cached_context.phase
phase = _advance_gait(env, context)
_GAIT_PHASE_CACHE[env] = (context, counter)
return phase
def _resample_gait(env: _G1Env, context: _G1GaitContext, env_ids: torch.Tensor) -> None:
device = context.phase.device
rows = env_ids.to(dtype=torch.int64, device=device)
count = rows.numel()
if count == 0:
return
rng_owner = getattr(env, "rng", None)
if not isinstance(rng_owner, ManagerRng):
raise RuntimeError("G1 gait phase reset requires the Manager-owned Torch generator")
left = rng_owner.uniform(0.0, 2.0 * math.pi, (count,), dtype=torch.float32)
if context.init_mode == "independent":
right = rng_owner.uniform(0.0, 2.0 * math.pi, (count,), dtype=torch.float32)
else:
right = left + torch.pi
samples = torch.stack((left, right), dim=1)
context.phase.index_copy_(0, rows, samples.to(dtype=context.phase.dtype))
def _resample_gait_mask(env: _G1Env, context: _G1GaitContext, reset_mask: torch.Tensor) -> None:
rng_owner = getattr(env, "rng", None)
if not isinstance(rng_owner, ManagerRng):
raise RuntimeError("G1 gait phase reset requires the Manager-owned Torch generator")
count = context.phase.shape[0]
left = rng_owner.uniform(0.0, 2.0 * math.pi, (count,), dtype=torch.float32)
if context.init_mode == "independent":
right = rng_owner.uniform(0.0, 2.0 * math.pi, (count,), dtype=torch.float32)
else:
right = left + torch.pi
samples = torch.stack((left, right), dim=1)
context.phase.copy_(
torch.where(reset_mask[:, None], samples.to(dtype=context.phase.dtype), context.phase)
)
[docs]
class G1GaitPhase(ManagerTermBase):
"""Two-foot gait phase observation in radians, with per-reset phase sampling."""
_allowed_params: ClassVar[frozenset[str]] = frozenset({"frequency", "init_mode"})
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(env)
unexpected = set(cfg.params) - self._allowed_params
if unexpected:
raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}")
frequency = _real(self.name, "frequency", cfg.params.get("frequency", 1.5), minimum=0.0)
init_mode = cfg.params.get("init_mode", "offset_phase")
if init_mode not in _GAIT_INIT_MODES:
raise ValueError(f"{self.name} init_mode must be one of {_GAIT_INIT_MODES}")
self._context = _gait_context(env, self.name, frequency, init_mode)
self._reset_mask_buffer = ResetMaskBuffer(env.num_envs, torch.device(env.device))
[docs]
def __call__(self, env: _G1Env, **params: Any) -> torch.Tensor:
del params
phase = _advance_gait(env, self._context)
if isinstance(phase, torch.Tensor):
return phase
return phase.copy()
[docs]
def reset(self, env_ids: torch.Tensor) -> None:
env = cast("_G1Env", self._env)
_resample_gait_mask(env, self._context, self._reset_mask_buffer.scatter(env_ids))
class _GaitRewardTerm(_SensorTerm):
"""Shared gait-context, forward-speed gate, and foot-position binding."""
_allowed_params = frozenset(
{
"frequency",
"init_mode",
"swing_height",
"tracking_sigma",
"min_forward_speed",
"command_name",
}
)
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(cfg, env)
frequency = _real(self.name, "frequency", cfg.params.get("frequency", 1.5), minimum=0.0)
init_mode = cfg.params.get("init_mode", "offset_phase")
if init_mode not in _GAIT_INIT_MODES:
raise ValueError(f"{self.name} init_mode must be one of {_GAIT_INIT_MODES}")
self._swing_height = _real(
self.name, "swing_height", cfg.params.get("swing_height", 0.09), minimum=0.0
)
self._tracking_sigma = _real(
self.name,
"tracking_sigma",
cfg.params.get("tracking_sigma", 0.008),
minimum=0.0,
strict_minimum=True,
)
self._min_forward_speed = _real(
self.name,
"min_forward_speed",
cfg.params.get("min_forward_speed", 0.0),
minimum=0.0,
)
command_name = cfg.params.get("command_name", "twist")
if not isinstance(command_name, str) or not command_name:
raise ValueError(f"{self.name} command_name must be a non-empty string")
self._command_name = command_name
self._context = _gait_context(env, self.name, frequency, init_mode)
self._feet_pos = self._bind(_FOOT_POS_SENSORS)
if self._feet_pos.dimensions != (3, 3):
raise ValueError(
f"{self.name} foot position sensors must expose 3-D xyz; received "
f"{self._feet_pos.dimensions} on backend '{self._feet_pos.backend_type}'"
)
self._linvel = self._bind(("pelvis_local_linvel",))
self._phase_inputs_view = self._bind((*_FOOT_POS_SENSORS, "pelvis_local_linvel"))
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return (*_FOOT_POS_SENSORS, "pelvis_local_linvel")
def _targets(self, env: _G1Env) -> tuple[torch.Tensor, torch.Tensor]:
phase = _cached_gait_phase(env, self._context)
return compute_feet_phase_height_targets(phase, self._swing_height)
def _foot_heights(self, env: _G1Env) -> tuple[torch.Tensor, torch.Tensor]:
values = self._read_state_tensor(
env, "foot position", (self._feet_pos, _FOOT_POS_SENSORS), (self.num_envs, 6)
)
return values[:, 2], values[:, 5]
def _gate(self, env: ManagerBasedRlEnv) -> torch.Tensor:
linvel = self._read_state_tensor(
env,
"pelvis linear velocity",
(self._linvel, ("pelvis_local_linvel",)),
(env.num_envs, 3),
)
forward_speed = torch.clamp(_torch_values(linvel)[:, 0], min=0.0)
return (forward_speed >= self._min_forward_speed).to(dtype=torch.float32)
def _gate_from_forward_speed(self, forward_speed: torch.Tensor) -> torch.Tensor:
positive_speed = torch.clamp(_torch_values(forward_speed), min=0.0)
return (positive_speed >= self._min_forward_speed).to(dtype=torch.float32)
def _phase_reward_inputs(
self, env: _G1Env
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Read foot height, forward speed, and both height targets in one call."""
values = self._read_state_tensor(
env,
"foot position and pelvis linear velocity",
(self._phase_inputs_view, (*_FOOT_POS_SENSORS, "pelvis_local_linvel")),
(self.num_envs, 9),
)
phase = _cached_gait_phase(env, self._context)
left_target, right_target = compute_feet_phase_height_targets(phase, self._swing_height)
return values[:, 2], values[:, 5], values[:, 6], left_target, right_target
[docs]
class feet_phase(_GaitRewardTerm):
"""Reward gait phase tracking by encouraging the expected swing-foot height."""
[docs]
def __call__(self, env: _G1Env, **params: Any) -> torch.Tensor:
del params
left_height, right_height, forward_speed, left_target, right_target = (
self._phase_reward_inputs(env)
)
error = _square(left_height - left_target) + _square(right_height - right_target)
reward = _exp(-error / self._tracking_sigma)
return reward * self._gate_from_forward_speed(forward_speed)
[docs]
class feet_phase_contrast(_GaitRewardTerm):
"""Reward left/right foot-height contrast against the gait-phase targets."""
[docs]
def __call__(self, env: _G1Env, **params: Any) -> torch.Tensor:
del params
left_target, right_target = self._targets(env)
left_height, right_height = self._foot_heights(env)
error = _square((left_height - right_height) - (left_target - right_target))
reward = _exp(-error / self._tracking_sigma)
return reward * self._gate(env)
class _FootContactTerm(_GaitRewardTerm):
"""Adds the aggregated per-foot contact binding shared by contact gait terms."""
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(cfg, env)
self._contacts = self._bind(_FOOT_CONTACT_SENSORS)
if self._contacts.dimensions != (1,) * 8:
raise ValueError(
f"{self.name} foot contact sensors must each expose 1-D found; received "
f"{self._contacts.dimensions} on backend '{self._contacts.backend_type}'"
)
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return (*super().tensor_sensor_names, *_FOOT_CONTACT_SENSORS)
def _contact_pair(self, env: ManagerBasedRlEnv) -> tuple[torch.Tensor, torch.Tensor]:
values = self._read_state_tensor(
env, "foot contact", (self._contacts, _FOOT_CONTACT_SENSORS), (self.num_envs, 8)
)
contacts = _torch_values(values) > 0.5
return torch.any(contacts[:, :4], dim=1), torch.any(contacts[:, 4:], dim=1)
[docs]
class feet_double_stance(_FootContactTerm):
"""Penalize double-stance contact while a forward command is active."""
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
del params
command = _command(env, self.name, self._command_name)
left_contact, right_contact = self._contact_pair(env)
command_tensor = _torch_values(command)
left = left_contact.to(device=command_tensor.device)
right = right_contact.to(device=command_tensor.device)
double_stance = (left & right).to(dtype=torch.float32)
forward_mask = (command_tensor[:, 0].clamp(min=0.0) > 1.0e-6).to(dtype=torch.float32)
return double_stance * forward_mask
[docs]
class feet_air_time(_FootContactTerm):
"""Count feet whose current air time sits inside the rewarded window."""
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(cfg, env)
self._air_time = torch.zeros(
(env.num_envs, 2), dtype=torch.float32, device=torch.device(env.device)
)
[docs]
def reset(self, env_ids: torch.Tensor) -> None:
self._air_time[env_ids] = 0.0
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
del params
left_contact, right_contact = self._contact_pair(env)
contact = torch.stack(
(
left_contact.to(device=self._air_time.device),
right_contact.to(device=self._air_time.device),
),
dim=1,
)
step_dt = _real(self.name, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True)
self._air_time = torch.where(
contact,
torch.zeros_like(self._air_time),
self._air_time + step_dt,
)
in_range = (self._air_time > 0.05) & (self._air_time < 0.5)
return in_range.to(dtype=torch.float32).sum(dim=1)
# ---------------------------------------------------------------------------
# Velocity / orientation reward terms (sensor-bound, G1Walk family equations)
# ---------------------------------------------------------------------------
class _LinVelTerm(_SensorTerm):
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self._linvel = self._bind(("pelvis_local_linvel",))
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return ("pelvis_local_linvel",)
def _read_linvel(self, env: ManagerBasedRlEnv) -> torch.Tensor:
return self._read_state_tensor(
env,
"pelvis linear velocity",
(self._linvel, ("pelvis_local_linvel",)),
(env.num_envs, 3),
)
class _GyroTerm(_SensorTerm):
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self._gyro = self._bind(("torso_gyro",))
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return ("torso_gyro",)
def _read_gyro(self, env: ManagerBasedRlEnv) -> torch.Tensor:
return self._read_state_tensor(
env, "torso gyro", (self._gyro, ("torso_gyro",)), (env.num_envs, 3)
)
class _UpvectorTerm(_SensorTerm):
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self._upvector = self._bind(("torso_upvector",))
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return ("torso_upvector",)
def _read_upvector(self, env: ManagerBasedRlEnv) -> torch.Tensor:
return self._read_state_tensor(
env, "torso upvector", (self._upvector, ("torso_upvector",)), (env.num_envs, 3)
)
[docs]
class forward_progress(_LinVelTerm):
"""Reward forward progress relative to commanded speed."""
_allowed_params = frozenset({"command_name"})
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
command = _command(env, self.name, _term_command_name(self.name, params))
linvel = self._read_linvel(env)
command = _torch_values(command)
linvel = _torch_values(linvel).to(device=command.device)
commanded_speed = torch.clamp(command[:, 0], min=1e-6)
forward_speed = torch.clamp(linvel[:, 0], min=0.0)
return torch.clamp(forward_speed / commanded_speed, max=1.0).to(dtype=torch.float32)
def _term_command_name(term: str, params: dict[str, Any]) -> str:
name = params.get("command_name", "twist")
if not isinstance(name, str) or not name:
raise ValueError(f"{term} command_name must be a non-empty string")
return name
[docs]
class under_speed(_LinVelTerm):
"""Penalty for being below commanded forward speed."""
_allowed_params = frozenset({"command_name"})
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
command = _command(env, self.name, _term_command_name(self.name, params))
linvel = self._read_linvel(env)
command = _torch_values(command)
linvel = _torch_values(linvel).to(device=command.device)
commanded_speed = torch.clamp(command[:, 0], min=1e-6)
forward_speed = torch.clamp(linvel[:, 0], min=0.0)
gap = torch.clamp(command[:, 0] - forward_speed, min=0.0)
return (gap / commanded_speed).to(dtype=torch.float32)
[docs]
class g1_tilt_exceeded(_UpvectorTerm):
"""Terminate when the base tilt from upright exceeds ``max_tilt_deg``."""
_allowed_params = frozenset({"max_tilt_deg"})
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
max_tilt_deg = _real(self.name, "max_tilt_deg", cfg.params.get("max_tilt_deg"), minimum=0.0)
self._max_tilt_rad = math.radians(max_tilt_deg)
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
del params
upvector = self._read_upvector(env)
z = torch.clamp(_torch_values(upvector)[:, 2], min=-1.0, max=1.0)
return torch.acos(z) > self._max_tilt_rad
[docs]
class g1_base_height_below_minimum(ManagerTermBase):
"""Terminate when root height drops below the task threshold.
Declaring the root body lets the packed device-resident read phase serve
this term without the retired NumPy body-state facade or a full host cache
refresh.
"""
_allowed_params = frozenset({"minimum_height"})
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(env)
self._minimum_height = _real(self.name, "minimum_height", cfg.params.get("minimum_height"))
asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)
if not isinstance(asset_cfg, SceneEntityCfg):
raise TypeError(f"{self.name} asset_cfg must be a SceneEntityCfg")
self._entity_name = asset_cfg.name
entity_cfg = env.cfg.scene.entities.get(self._entity_name) if env.cfg.scene else None
root = getattr(entity_cfg, "root_body_name", None) if entity_cfg is not None else None
if not isinstance(root, str) or not root:
raise ValueError(f"{self.name} requires an entity root_body_name")
self._root_body_name = root
@property
def tensor_body_names(self) -> tuple[str, ...]:
return (self._root_body_name,)
@property
def entity_name(self) -> str:
return self._entity_name
[docs]
def __call__(self, env: _G1Env, **params: Any) -> torch.Tensor:
del params
read_plan = getattr(env.scene, "read_plan", None)
if read_plan is None:
raise NotImplementedError(
f"Manager term '{self.name}' requires the packed root-body read plan"
)
body = read_plan.body_tensor_view(self._entity_name, (self._root_body_name,))
return body.pos_w[:, 0, 2] < self._minimum_height
[docs]
class penalty_feet_ori(_SensorTerm):
"""Penalty for non-flat foot orientations (roll/pitch quaternion rows)."""
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self._quats = self._bind(_FOOT_QUAT_SENSORS)
if self._quats.dimensions != (4, 4):
raise ValueError(
f"{self.name} foot quaternion sensors must expose 4-D quats; received "
f"{self._quats.dimensions} on backend '{self._quats.backend_type}'"
)
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return _FOOT_QUAT_SENSORS
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
del params
values = self._read_state_tensor(
env,
"foot quaternion",
(self._quats, _FOOT_QUAT_SENSORS),
(env.num_envs, 8),
)
penalty = _square(values[:, 1]) + _square(values[:, 2])
penalty = penalty + _square(values[:, 5]) + _square(values[:, 6])
return penalty
[docs]
class penalty_close_feet_xy(_SensorTerm):
"""Penalty for feet closer than ``threshold`` in the horizontal plane."""
_allowed_params = frozenset({"threshold"})
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
self._threshold = _real(
self.name, "threshold", cfg.params.get("threshold", 0.15), minimum=0.0
)
self._feet_pos = self._bind(_FOOT_POS_SENSORS)
if self._feet_pos.dimensions != (3, 3):
raise ValueError(
f"{self.name} foot position sensors must expose 3-D xyz; received "
f"{self._feet_pos.dimensions} on backend '{self._feet_pos.backend_type}'"
)
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return _FOOT_POS_SENSORS
[docs]
def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> torch.Tensor:
del params
values = self._read_state_tensor(
env,
"foot position",
(self._feet_pos, _FOOT_POS_SENSORS),
(env.num_envs, 6),
)
positions = _torch_values(values)
feet_delta = positions[:, :2] - positions[:, 3:5]
feet_dist = torch.linalg.vector_norm(feet_delta, dim=1)
feet_gap = torch.square(feet_dist - self._threshold)
return torch.where(
feet_dist < self._threshold,
feet_gap,
torch.zeros_like(feet_dist),
).to(dtype=torch.float32)
# ---------------------------------------------------------------------------
# Entity-based reward terms (G1Walk family equations, no sensor binding)
# ---------------------------------------------------------------------------
[docs]
class weighted_pose(ManagerTermBase):
"""Weighted L2 penalty for joint position deviation from the default pose."""
[docs]
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
super().__init__(env)
asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)
if not isinstance(asset_cfg, SceneEntityCfg):
raise TypeError(f"{self.name} asset_cfg must be a SceneEntityCfg")
self._entity_name = asset_cfg.name
self._joint_ids = asset_cfg.joint_ids
self._weights = _weights(self.name, "pose_weights", cfg.params["pose_weights"]).to(
device=env.device
)
@property
def tensor_state_entity(self) -> str:
return self._entity_name
[docs]
def __call__(
self,
env: ManagerBasedRlEnv,
pose_weights: Any,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
del pose_weights, asset_cfg
read_plan = getattr(env.scene, "read_plan", None)
if read_plan is None:
raise NotImplementedError(
f"{self.name} requires the Manager packed joint-state read plan"
)
position = read_plan.joint_tensor_view(self._entity_name).joint_pos[:, self._joint_ids]
default = env.scene[self._entity_name].data.default_joint_pos(env.device)[
:, self._joint_ids
]
if self._weights.numel() != position.shape[1]:
raise ValueError(
f"{self.name} pose_weights length {self._weights.numel()} does not match "
f"joint count {position.shape[1]}"
)
return (self._weights * (position - default).square()).sum(dim=1)
[docs]
def reset(self, env_ids: torch.Tensor) -> None:
del env_ids
[docs]
class upper_body_pose(weighted_pose):
"""Weighted L2 pose penalty with the twelve leg joints zeroed out."""
[docs]
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
weights = _weights(self.name, "pose_weights", cfg.params["pose_weights"])
if weights.numel() < 12:
raise ValueError(f"{self.name} pose_weights must cover at least the twelve leg joints")
weights = weights.clone()
weights[:12] = 0.0
cfg.params["pose_weights"] = weights
super().__init__(cfg, env)
# ---------------------------------------------------------------------------
# Velocity command term (G1Walk dead zone + standing zeroing semantics)
# ---------------------------------------------------------------------------
[docs]
@dataclass(kw_only=True)
class G1VelocityCommandCfg(UniformVelocityCommandCfg):
"""G1 velocity command with the G1Walk-family planar-norm dead zone."""
planar_dead_zone: float = 0.2
[docs]
def build(self, env: ManagerBasedRlEnv) -> G1VelocityCommand:
return G1VelocityCommand(self, env)
[docs]
class G1VelocityCommand(UniformVelocityCommand):
cfg: G1VelocityCommandCfg # pyright: ignore[reportIncompatibleVariableOverride]
[docs]
def __init__(self, cfg: G1VelocityCommandCfg, env: ManagerBasedRlEnv):
self._planar_dead_zone = _real(
"G1VelocityCommand", "planar_dead_zone", cfg.planar_dead_zone, minimum=0.0
)
if cfg.heading_command:
raise NotImplementedError(
"G1VelocityCommand capability 'heading command' is unavailable in the "
"Manager-Based runtime; the G1Walk heading-channel zeroing was not "
"ported and fails closed instead of falling back"
)
super().__init__(cfg, env)
def _resample_command(self, reset_mask: torch.Tensor) -> None:
super()._resample_command(reset_mask)
tensor = self._tensor_command
assert tensor is not None
planar = tensor[:, :2]
moving = torch.linalg.vector_norm(planar, dim=1) > self._planar_dead_zone
tensor[:, :2] = torch.where(reset_mask[:, None], planar * moving[:, None], planar)
def _merge_samples_masked(self, samples: torch.Tensor, reset_mask: torch.Tensor) -> None:
moving = torch.linalg.vector_norm(samples[:, :2], dim=1) > self._planar_dead_zone
samples = samples.clone()
samples[:, :2] *= moving[:, None]
super()._merge_samples_masked(samples, reset_mask)
# ---------------------------------------------------------------------------
# Penalty curriculum (ports the G1Walk family penalty-curriculum semantics)
# ---------------------------------------------------------------------------
[docs]
class G1PenaltyCurriculum(ManagerTermBase):
"""Scale negative-weight reward terms by average episode length.
Ports the G1Walk task-family penalty curriculum (64c7f8b8, roadmap #1042):
penalty weights start at ``initial_scale`` of their configured value and
relax toward ``max_scale`` as the tracked average episode length crosses
the configured thresholds.
"""
_allowed_params: ClassVar[frozenset[str]] = frozenset(
{
"initial_scale",
"min_scale",
"max_scale",
"level_down_threshold",
"level_up_threshold",
"degree",
}
)
[docs]
def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env):
super().__init__(env)
unexpected = set(cfg.params) - self._allowed_params
if unexpected:
raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}")
self._min_scale = _real(
self.name, "min_scale", cfg.params.get("min_scale", 0.5), minimum=0.0
)
self._max_scale = _real(
self.name,
"max_scale",
cfg.params.get("max_scale", 1.0),
minimum=self._min_scale,
)
self._current_scale = _real(
self.name,
"initial_scale",
cfg.params.get("initial_scale", 0.5),
minimum=self._min_scale,
)
if self._current_scale > self._max_scale:
raise ValueError(f"{self.name} initial_scale must be at most max_scale")
self._level_down_threshold = _real(
self.name,
"level_down_threshold",
cfg.params.get("level_down_threshold", 150.0),
minimum=0.0,
)
self._level_up_threshold = _real(
self.name,
"level_up_threshold",
cfg.params.get("level_up_threshold", 750.0),
minimum=0.0,
)
self._degree = _real(self.name, "degree", cfg.params.get("degree", 0.001), minimum=0.0)
self._tracker = EpisodeLengthTracker(env.num_envs, device=env.device)
self._reset_mask_buffer = ResetMaskBuffer(env.num_envs, torch.device(env.device))
self._original_weights: dict[str, float] = {}
for name in env.reward_manager.active_terms:
weight = float(env.reward_manager.get_term_cfg(name).weight)
if weight < 0.0:
self._original_weights[name] = weight
self._apply_scale()
def _apply_scale(self) -> None:
env = cast("_G1Env", self._env)
for name, original in self._original_weights.items():
env.reward_manager.get_term_cfg(name).weight = original * self._current_scale
[docs]
def __call__(
self,
env: _G1Env,
env_ids: torch.Tensor | slice | None,
**params: Any,
) -> dict[str, float]:
del params
device = env.device
if env_ids is None:
ids = torch.arange(env.num_envs, dtype=torch.int64, device=device)
elif isinstance(env_ids, slice):
ids = torch.arange(env.num_envs, dtype=torch.int64, device=device)[env_ids]
elif isinstance(env_ids, torch.Tensor):
ids = env_ids.to(device=device, dtype=torch.int64).reshape(-1)
else:
raise TypeError(
f"{self.name} env_ids must be Torch tensor, slice, or None; "
f"got {type(env_ids).__name__}"
)
reset_buf = env.reset_buf
if not isinstance(reset_buf, torch.Tensor):
raise TypeError(f"{self.name} requires reset_buf to be a Torch tensor")
done_ids = ids[reset_buf.index_select(0, ids)]
episode_lengths = env.episode_length_buf
if not isinstance(episode_lengths, torch.Tensor):
raise TypeError(f"{self.name} requires episode_length_buf to be a Torch tensor")
if done_ids.numel() > 0:
self._tracker.update(episode_lengths.index_select(0, done_ids))
average = self._tracker.average
# Reward weights are public Python floats, so this is the unavoidable
# scalar publication boundary. It happens once per curriculum update,
# not once per tracker batch.
average_value = float(average.item())
changed = False
if done_ids.numel() > 0:
if average_value < self._level_down_threshold:
self._current_scale *= 1.0 - self._degree
changed = True
elif average_value > self._level_up_threshold:
self._current_scale *= 1.0 + self._degree
changed = True
if changed:
self._current_scale = min(max(self._current_scale, self._min_scale), self._max_scale)
self._apply_scale()
return {
"average_episode_length": average_value,
"penalty_scale": self._current_scale,
}
[docs]
def reset(self, env_ids: torch.Tensor) -> None:
"""Row boundary for the global curriculum state."""
del env_ids
[docs]
def reset_scoped(
self,
env: _G1Env,
env_ids: torch.Tensor,
**params: Any,
) -> dict[str, float]:
"""Selected-row curriculum update at reset."""
del params
reset_mask = self._reset_mask_buffer.scatter(env_ids)
reset_buf = env.reset_buf
if not isinstance(reset_buf, torch.Tensor):
raise TypeError(f"{self.name} requires reset_buf to be a Torch tensor")
done_mask = reset_buf & reset_mask
episode_lengths = env.episode_length_buf
if not isinstance(episode_lengths, torch.Tensor):
raise TypeError(f"{self.name} requires episode_length_buf to be a Torch tensor")
done_count = done_mask.sum()
self._tracker.update_masked(episode_lengths, done_mask)
average = self._tracker.average
average_value = float(average.item())
changed = False
if bool(done_count.item() > 0):
if average_value < self._level_down_threshold:
self._current_scale *= 1.0 - self._degree
changed = True
elif average_value > self._level_up_threshold:
self._current_scale *= 1.0 + self._degree
changed = True
if changed:
self._current_scale = min(max(self._current_scale, self._min_scale), self._max_scale)
self._apply_scale()
return {
"average_episode_length": average_value,
"penalty_scale": self._current_scale,
}
# ---------------------------------------------------------------------------
# G1 Manager-Based env: Registry-owned production runtime on the single
# lifecycle
# ---------------------------------------------------------------------------
[docs]
@dataclass(kw_only=True)
class G1WalkResetOwnerCfg(ResetOwnerCfg):
"""Fuse the selected-reset lifecycle for the canonical G1 walk command."""
command_name: str = "twist"
owns_observation_reset: bool = False
[docs]
class G1WalkResetOwner(ResetOwner):
"""Own the tensor G1 command reset and skip redundant host validations.
The walk command is sampled from bounded Torch uniforms and the G1 dead-zone
write is bounded by construction. Re-running the generic full-command
finite synchronization after that device-resident transaction only adds a
device synchronization on the reset hot path.
"""
[docs]
def __init__(self, cfg: G1WalkResetOwnerCfg, env: _G1Env):
super().__init__(cfg, env)
command_manager = getattr(env, "command_manager", None)
if command_manager is None:
raise KeyError("G1WalkResetOwner requires a configured command manager")
try:
command = command_manager.get_term(cfg.command_name)
except (AttributeError, KeyError) as exc:
raise KeyError(
f"G1WalkResetOwner command term '{cfg.command_name}' is not configured"
) from exc
if not isinstance(command, UniformVelocityCommand):
raise TypeError("G1WalkResetOwner requires a UniformVelocityCommand term")
if getattr(env, "rng", None) is None:
raise NotImplementedError("G1WalkResetOwner requires the Manager RNG")
[docs]
def reset_transaction(self, env_ids: torch.Tensor) -> dict[str, float]:
env = cast("_G1Env", self._env)
command = env.command_manager.get_term(self.cfg.command_name)
if not callable(getattr(command, "reset", None)):
raise NotImplementedError(
f"G1WalkResetOwner command '{self.cfg.command_name}' does not define "
"the reset boundary"
)
command.reset(env_ids)
return {}
[docs]
def reset_committed(self, env_ids: torch.Tensor) -> dict[str, float]:
env = cast("_G1Env", self._env)
env.action_manager.clear_action_state(env_ids)
env.metrics_manager.clear_episode_state(env_ids)
return {}
[docs]
@dataclass(kw_only=True)
class G1WalkRewardPackCfg(RewardTermCfg):
"""Fused canonical G1 walk reward family for tensor-resident owners."""
command_name: str = "twist"
entity_name: str = "robot"
tracking_lin_sigma: float = 0.25
tracking_ang_sigma: float = 0.25
gait_frequency: float = 1.5
gait_init_mode: str = "offset_phase"
swing_height: float = 0.09
feet_phase_sigma: float = 0.04
min_forward_speed: float = 0.0
pose_weights: tuple[float, ...] = tuple(1.0 for _ in range(29))
_G1_WALK_REWARD_PACK_PARAMS: frozenset[str] = frozenset(
{
"command_name",
"entity_name",
"tracking_lin_sigma",
"tracking_ang_sigma",
"gait_frequency",
"gait_init_mode",
"swing_height",
"feet_phase_sigma",
"min_forward_speed",
"pose_weights",
}
)
[docs]
class G1WalkRewardPack(ManagerTermBase):
"""Evaluate the canonical walk reward family in one tensor carrier read.
This pack fuses the G1 owner's state-only per-step rewards: linear/angular
tracking, angular-xy and orientation penalties, action rate, weighted pose,
foot-orientation penalty, and phase tracking. Equations, weights, and
declaration order remain owned by ``RewardManager``; this owner only removes
duplicate carrier reads and intermediate per-term kernel launches.
"""
returns_transient_tensor = True
last_step_timing_ms: dict[str, float] = {}
[docs]
def __init__(self, cfg: G1WalkRewardPackCfg, env: _G1Env):
super().__init__(env)
if not isinstance(cfg, G1WalkRewardPackCfg):
raise TypeError("G1WalkRewardPack requires G1WalkRewardPackCfg")
unexpected = set(cfg.params) - _G1_WALK_REWARD_PACK_PARAMS
if unexpected:
raise TypeError(
f"G1WalkRewardPack received unsupported parameters: {sorted(unexpected)}"
)
self.cfg = cfg
self._device = torch.device(env.device)
self._command = _command(env, self.name, cfg.command_name)
self._gait_context = _gait_context(env, self.name, cfg.gait_frequency, cfg.gait_init_mode)
self._entity = cast("Entity", env.scene[cfg.entity_name])
self._pose_weights = torch.as_tensor(
cfg.pose_weights, dtype=torch.float32, device=self._device
)
if self._pose_weights.ndim != 1 or self._pose_weights.numel() == 0:
raise ValueError("G1WalkRewardPack pose_weights must be a non-empty vector")
self._tracking_lin_scale = cfg.tracking_lin_sigma * cfg.tracking_lin_sigma
self._tracking_ang_scale = cfg.tracking_ang_sigma * cfg.tracking_ang_sigma
self._feet_phase_scale = cfg.feet_phase_sigma
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return (
"pelvis_local_linvel",
"torso_gyro",
"torso_upvector",
*_FOOT_QUAT_SENSORS,
*_FOOT_POS_SENSORS,
)
def _read_sensors(self, env: _G1Env) -> dict[str, torch.Tensor]:
read_plan = getattr(env.scene, "read_plan", None)
if read_plan is None or read_plan.device.type != self._device.type:
raise NotImplementedError(
"G1WalkRewardPack requires the declared tensor sensor read plan"
)
names = self.tensor_sensor_names
owner_sensors = set(read_plan.sensor_names.get(self.cfg.entity_name, ()))
if not set(names).issubset(owner_sensors):
missing = sorted(set(names) - owner_sensors)
raise NotImplementedError(f"G1WalkRewardPack requires device tensor sensors {missing}")
entity = env.scene[self.cfg.entity_name]
return cast(dict[str, torch.Tensor], read_plan.sensor_tensor_views(entity, names).values)
[docs]
def __call__(self, env: _G1Env) -> torch.Tensor:
call_started = time.perf_counter()
values = self._read_sensors(env)
command = self._command
if not isinstance(command, torch.Tensor):
raise TypeError("G1WalkRewardPack requires a tensor command carrier")
linvel = values["pelvis_local_linvel"]
gyro = values["torso_gyro"]
upvector = values["torso_upvector"]
foot_quat = values[_FOOT_QUAT_SENSORS[0]]
for name in _FOOT_QUAT_SENSORS[1:]:
foot_quat = torch.cat((foot_quat, values[name]), dim=1)
foot_pos = values[_FOOT_POS_SENSORS[0]]
for name in _FOOT_POS_SENSORS[1:]:
foot_pos = torch.cat((foot_pos, values[name]), dim=1)
action_manager = cast(Any, self._env.action_manager)
joint_state = getattr(env.scene, "read_plan", None)
if joint_state is None:
raise NotImplementedError("G1WalkRewardPack requires the tensor joint read plan")
joint_pos = joint_state.joint_tensor_view(self._entity).joint_pos
default = self._entity.data.default_joint_pos(self._device)
phase = _cached_gait_phase(env, self._gait_context)
targets = compute_feet_phase_height_targets(phase, self.cfg.swing_height)
left_target = cast(torch.Tensor, targets[0])
right_target = cast(torch.Tensor, targets[1])
lin_error = torch.sum(torch.square(command[:, :2] - linvel[:, :2]), dim=1)
ang_error = torch.square(command[:, 2] - gyro[:, 2])
ang_xy_error = torch.sum(torch.square(gyro[:, :2]), dim=1)
ori_error = torch.sum(torch.square(upvector[:, :2]), dim=1)
action_delta = action_manager.action - action_manager.prev_action
action_error = torch.sum(torch.square(action_delta), dim=1)
pose_error = torch.sum(self._pose_weights * torch.square(joint_pos - default), dim=1)
foot_ori_error = (
torch.square(foot_quat[:, 1])
+ torch.square(foot_quat[:, 2])
+ torch.square(foot_quat[:, 5])
+ torch.square(foot_quat[:, 6])
)
phase_error = torch.square(foot_pos[:, 2] - left_target) + torch.square(
foot_pos[:, 5] - right_target
)
forward_speed = linvel[:, 0]
gate = (torch.clamp(forward_speed, min=0.0) >= self.cfg.min_forward_speed).to(
dtype=torch.float32
)
outputs = [
torch.exp(-lin_error / self._tracking_lin_scale),
torch.exp(-ang_error / self._tracking_ang_scale),
ang_xy_error,
ori_error,
action_error,
pose_error,
foot_ori_error,
torch.exp(-phase_error / self._feet_phase_scale) * gate,
torch.ones_like(lin_error),
]
result = torch.stack(outputs, dim=1)
self.last_step_timing_ms["update_state_reward_g1_walk_pack_call_ms"] = (
time.perf_counter() - call_started
) * 1000.0
return result
[docs]
@dataclass(kw_only=True)
class G1WalkTerminationPackCfg(TerminationTermCfg):
"""Fused canonical G1 walk timeout, tilt, and height termination."""
entity_name: str = "robot"
root_body_name: str = "pelvis"
max_tilt_deg: float = 65.0
minimum_height: float = 0.3
[docs]
class G1WalkTerminationPack(ManagerTermBase):
"""Evaluate the canonical G1 walk failure terms in one carrier read."""
returns_transient_tensor = True
[docs]
def __init__(self, cfg: G1WalkTerminationPackCfg, env: _G1Env):
super().__init__(env)
if not isinstance(cfg, G1WalkTerminationPackCfg):
raise TypeError("G1WalkTerminationPack requires G1WalkTerminationPackCfg")
unexpected = set(cfg.params) - {
"entity_name",
"root_body_name",
"max_tilt_deg",
"minimum_height",
}
if unexpected:
raise TypeError(
f"G1WalkTerminationPack received unsupported parameters: {sorted(unexpected)}"
)
self.cfg = cfg
self._device = torch.device(env.device)
self._entity = cast("Entity", env.scene[cfg.entity_name])
self._max_tilt_rad = math.radians(
_real(
"G1WalkTerminationPack",
"max_tilt_deg",
cfg.max_tilt_deg,
minimum=0.0,
)
)
self._minimum_height = _real(
"G1WalkTerminationPack",
"minimum_height",
cfg.minimum_height,
)
self._tilt_threshold = math.cos(self._max_tilt_rad)
@property
def tensor_sensor_names(self) -> tuple[str, ...]:
return ("torso_upvector",)
@property
def tensor_body_names(self) -> tuple[str, ...]:
return (self.cfg.root_body_name,)
@property
def entity_name(self) -> str:
return self.cfg.entity_name
[docs]
def reset(self, env_ids: torch.Tensor) -> None:
"""Row boundary for the stateless termination pack."""
del env_ids
[docs]
def __call__(self, env: _G1Env) -> torch.Tensor:
read_plan = getattr(env.scene, "read_plan", None)
if read_plan is None or read_plan.device.type != self._device.type:
raise NotImplementedError(
"G1WalkTerminationPack requires the declared tensor sensor/body read plan"
)
owner_sensors = set(read_plan.sensor_names.get(self.cfg.entity_name, ()))
if "torso_upvector" not in owner_sensors:
raise NotImplementedError(
"G1WalkTerminationPack requires device tensor sensor 'torso_upvector'"
)
entity = env.scene[self.cfg.entity_name]
upvector = read_plan.sensor_tensor_views(entity, ("torso_upvector",)).values[
"torso_upvector"
]
body = read_plan.body_tensor_view(self.cfg.entity_name, (self.cfg.root_body_name,))
return (upvector[:, 2] < self._tilt_threshold) | (
body.pos_w[:, 0, 2] < self._minimum_height
)
[docs]
class G1WalkManagerBasedEnv(_ConcreteManagerBasedRlEnv):
"""Manager-Based G1 walk runtime."""
[docs]
def make_g1_walk_env(
cfg: ManagerBasedRlEnvCfg,
num_envs: int = 1,
backend_type: str = "mujoco",
) -> G1WalkManagerBasedEnv:
"""Construct the Registry-owned G1 Manager-Based production runtime."""
if not isinstance(cfg, ManagerBasedRlEnvCfg):
raise TypeError(
f"make_g1_walk_env expected ManagerBasedRlEnvCfg, received {type(cfg).__name__}"
)
if isinstance(num_envs, bool) or not isinstance(num_envs, int) or num_envs <= 0:
raise ValueError(f"make_g1_walk_env num_envs must be a positive integer, got {num_envs!r}")
if not isinstance(backend_type, str) or not backend_type:
raise ValueError(
f"make_g1_walk_env backend_type must be a non-empty string, got {backend_type!r}"
)
cfg.validate()
assert cfg.scene is not None
base_name, body_state_requested, tracked_body_names = _resolve_backend_entity_contract(cfg)
backend_kwargs = env_backend_kwargs(cfg)
backend_kwargs["base_name"] = base_name
if backend_type == "mujoco" and tracked_body_names is not None:
backend_kwargs["tracked_body_names"] = tracked_body_names
backend = create_backend(
backend_type,
cfg.scene,
num_envs,
cfg.sim_dt,
body_state_required=body_state_requested,
**backend_kwargs,
)
try:
return G1WalkManagerBasedEnv(cfg, backend, num_envs)
except Exception:
backend.cleanup_scene_assets()
raise
__all__ = [
"G1GaitPhase",
"G1PenaltyCurriculum",
"G1VelocityCommand",
"G1VelocityCommandCfg",
"G1WalkResetOwner",
"G1WalkResetOwnerCfg",
"G1WalkManagerBasedEnv",
"G1WalkRewardPack",
"G1WalkTerminationPack",
"G1WalkTerminationPackCfg",
"G1WalkRewardPackCfg",
"compute_feet_phase_contact_targets",
"compute_feet_phase_height_targets",
"feet_air_time",
"feet_double_stance",
"feet_phase",
"feet_phase_contact",
"feet_phase_contrast",
"forward_progress",
"g1_tilt_exceeded",
"make_g1_walk_env",
"penalty_close_feet_xy",
"penalty_feet_ori",
"under_speed",
"upper_body_pose",
"weighted_pose",
]