53 lines
2.3 KiB
Python
53 lines
2.3 KiB
Python
"""Explicit straight/side/yaw command modes, with more stair traversal samples.
|
|
|
|
``deadzone`` zeroes per-axis commands below the threshold; v7 uses 0.05 so tiny
|
|
creep commands become standing instead of ambiguous near-zero motion.
|
|
"""
|
|
from dataclasses import dataclass
|
|
import torch
|
|
from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
|
|
|
|
|
|
def mix_commands(values, stairs):
|
|
result = values.clone()
|
|
mode = torch.rand(len(values),device=values.device)
|
|
# General terrain: 40% pure vx, 20% pure vy, 20% pure yaw, 20% combined.
|
|
result[mode<.4,1:] = 0.
|
|
result[(mode>=.4)&(mode<.6),0] = 0.
|
|
result[(mode>=.4)&(mode<.6),2] = 0.
|
|
result[(mode>=.6)&(mode<.8),:2] = 0.
|
|
# Stairs: 80% low-speed pure forward/backward, 20% full-range commands.
|
|
result[stairs] = values[stairs]
|
|
traverse = stairs & (mode<.8)
|
|
magnitude = .2 + .4*torch.rand(len(values),device=values.device)
|
|
result[traverse,0] = torch.where(values[traverse,0]<0,-magnitude[traverse],magnitude[traverse])
|
|
result[traverse,1:] = 0.
|
|
return result
|
|
|
|
|
|
class MixedVelocityCommand(UniformVelocityCommand):
|
|
def _resample_command(self, env_ids):
|
|
super()._resample_command(env_ids)
|
|
gen = self._env.cfg.scene.terrain.terrain_generator
|
|
stair_columns = torch.tensor([c.kind in ('up','down') for c in gen.sub_terrains.values()],device=self.device)
|
|
stairs = stair_columns[self._env.scene.terrain.terrain_types[env_ids]]
|
|
mixed = mix_commands(self.vel_command_b[env_ids],stairs)
|
|
# Stair magnitudes must obey the current curriculum bounds, including reverse.
|
|
for i, bounds in enumerate((self.cfg.ranges.lin_vel_x, self.cfg.ranges.lin_vel_y, self.cfg.ranges.ang_vel_z)):
|
|
mixed[:,i].clamp_(*bounds)
|
|
# Per-axis deadband: tiny commands are standing, not an ambiguous creep.
|
|
if self.cfg.deadzone > 0.:
|
|
mixed = torch.where(mixed.abs() < self.cfg.deadzone,
|
|
torch.zeros_like(mixed), mixed)
|
|
self.vel_command_b[env_ids] = mixed
|
|
if self.cfg.hold_stair_command:
|
|
self.time_left[env_ids[stairs]] = self._env.cfg.episode_length_s + 1.
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class MixedVelocityCommandCfg(UniformVelocityCommandCfg):
|
|
hold_stair_command: bool = False
|
|
deadzone: float = 0.
|
|
def build(self, env):
|
|
return MixedVelocityCommand(self,env)
|