Initial commit

This commit is contained in:
2026-11-02 17:10:39 +08:00
commit 102fdbaf47
79 changed files with 56509 additions and 0 deletions
+5
View File
@@ -0,0 +1,5 @@
"""S10 flat and terrain environment factories; no implicit task registration.
Use flat.s10_velocity_flat_env_cfg or
terrain.s10_velocity_terrain_env_cfg to construct an environment.
"""
+148
View File
@@ -0,0 +1,148 @@
"""Small flat-ground velocity task for the S10 mixed leg/wheel controller.
This environment is independent of the chase-specific SuperMimic learner.
"""
from copy import deepcopy
import torch
from mjlab.envs import mdp
from mjlab.envs.mdp.actions import JointPositionActionCfg, JointVelocityActionCfg
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
from mjlab.managers.reward_manager import RewardTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.tasks.velocity import mdp as velocity_mdp
from mjlab.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
from mjlab.utils.noise import UniformNoiseCfg
from s10_supermimic.robot.config import (
LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, WHEEL_GEOM_NAMES, get_s10_robot_cfg,
)
ACTOR_HISTORY = 5
def leg_posture_cost(env, asset_cfg: SceneEntityCfg) -> torch.Tensor:
d = env.scene[asset_cfg.name].data
return (d.joint_pos[:, asset_cfg.joint_ids]
- d.default_joint_pos[:, asset_cfg.joint_ids]).square().sum(-1)
def standstill_cost(env, wheel_cfg: SceneEntityCfg) -> torch.Tensor:
"""Extra braking at zero command, including wheels spinning in place.
This suppresses velocity; it does not fix a world-frame position or heading.
"""
command = env.command_manager.get_command("twist")
standing = command.square().sum(-1) < 1e-6
d = env.scene["robot"].data
cost = (d.root_link_lin_vel_b.square().sum(-1)
+ d.root_link_ang_vel_b.square().sum(-1)
+ 0.01 * d.joint_vel[:, wheel_cfg.joint_ids].square().sum(-1))
return standing * cost
def s10_velocity_flat_env_cfg(play: bool = False, history: int = ACTOR_HISTORY,
actor_wheel_vel: bool = True,
hipy: float | None = None,
hipx: float | None = None,
knee: float | None = None,
leg_kp: float = 80.0,
leg_kd: float = 2.0):
if history < 1:
raise ValueError("history must be at least 1")
cfg = make_velocity_env_cfg()
cfg.scene.entities = {"robot": get_s10_robot_cfg(
hipy=hipy, hipx=hipx, knee=knee, leg_kp=leg_kp, leg_kd=leg_kd)}
cfg.scene.terrain.terrain_type = "plane"
cfg.scene.terrain.terrain_generator = None
cfg.scene.sensors = ()
cfg.scene.num_envs = 1 # Training entry point must opt into larger batches.
cfg.sim.mujoco.cone = "elliptic"
cfg.sim.mujoco.impratio = 10
cfg.sim.mujoco.ccd_iterations = 50
cfg.viewer.body_name = "base_link"
cfg.viewer.distance = 2.0
legs = SceneEntityCfg("robot", joint_names=LEG_JOINT_NAMES, preserve_order=True)
wheels = SceneEntityCfg("robot", joint_names=WHEEL_JOINT_NAMES, preserve_order=True)
cfg.actions = {
# TODO(s10): flat pretraining keeps the smaller flat scales (0.25 / wheel 15).
# Terrain revision 4 overrides these to 0.25/0.6/0.8 and wheel 35. If you
# warm-start from a flat actor, consider matching the rev4 scales here.
"joint_pos": JointPositionActionCfg(
entity_name="robot", actuator_names=LEG_JOINT_NAMES,
scale=0.25, use_default_offset=True,
),
"wheel_vel": JointVelocityActionCfg(
entity_name="robot", actuator_names=WHEEL_JOINT_NAMES,
scale=15.0, use_default_offset=False, clip={".*": (-35.0, 35.0)},
),
}
# 53 values/frame (v4 / new v7, with measured wheel speed) or legacy 49
# values/frame (old v7). The critic and world model always keep wheel speed.
terms = {
"base_ang_vel": ObservationTermCfg(
func=mdp.base_ang_vel, noise=UniformNoiseCfg(n_min=-0.2, n_max=0.2)),
"projected_gravity": ObservationTermCfg(
func=mdp.projected_gravity, noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05)),
"joint_pos": ObservationTermCfg(
func=mdp.joint_pos_rel, params={"asset_cfg": legs, "biased": True},
noise=UniformNoiseCfg(n_min=-0.01, n_max=0.01)),
"joint_vel": ObservationTermCfg(
func=mdp.joint_vel_rel, params={"asset_cfg": legs}, scale=0.05,
noise=UniformNoiseCfg(n_min=-1.5, n_max=1.5)),
"wheel_vel": ObservationTermCfg(
func=mdp.joint_vel_rel, params={"asset_cfg": wheels}, scale=0.05,
noise=UniformNoiseCfg(n_min=-1.5, n_max=1.5)),
"actions": ObservationTermCfg(func=mdp.last_action),
"command": ObservationTermCfg(
func=mdp.generated_commands, params={"command_name": "twist"}),
}
critic = deepcopy(terms)
critic["joint_pos"].params["biased"] = False
critic["base_lin_vel"] = ObservationTermCfg(func=mdp.base_lin_vel)
if not actor_wheel_vel:
terms.pop("wheel_vel")
cfg.observations = {
"actor": ObservationGroupCfg(terms=terms, concatenate_terms=True,
enable_corruption=not play, history_length=history),
"critic": ObservationGroupCfg(terms=critic, concatenate_terms=True,
enable_corruption=False),
}
command = cfg.commands["twist"]
command.heading_command = False
command.ranges.heading = None
command.rel_heading_envs = 0.0
command.rel_standing_envs = 0.2
command.ranges.lin_vel_x = (-1.0, 1.0)
command.ranges.lin_vel_y = (-0.2, 0.2)
command.ranges.ang_vel_z = (-0.5, 0.5)
cfg.curriculum = {}
cfg.events.pop("push_robot", None)
cfg.events["base_com"].params["asset_cfg"].body_names = ("base_link",)
cfg.events["foot_friction"].params["asset_cfg"].geom_names = WHEEL_GEOM_NAMES
cfg.events["encoder_bias"].params["asset_cfg"] = legs
# Wheel rolling must not incur quadruped foot slip / swing height costs.
cfg.rewards = {
"track_linear_velocity": cfg.rewards["track_linear_velocity"],
"track_angular_velocity": cfg.rewards["track_angular_velocity"],
"upright": RewardTermCfg(func=velocity_mdp.upright, weight=1.0,
params={"std": 0.45, "asset_cfg": SceneEntityCfg(
"robot", body_names=("base_link",))}),
"leg_posture": RewardTermCfg(func=leg_posture_cost, weight=-0.2,
params={"asset_cfg": legs}),
"standstill": RewardTermCfg(func=standstill_cost, weight=-2.0,
params={"wheel_cfg": wheels}),
"action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.05),
"joint_torques_l2": RewardTermCfg(func=mdp.joint_torques_l2, weight=-2e-5),
"dof_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-1.0,
params={"asset_cfg": legs}),
"terminated": RewardTermCfg(func=mdp.is_terminated, weight=-20.0),
}
cfg.terminations.pop("out_of_terrain_bounds", None)
if play:
for event in ("base_com", "encoder_bias", "foot_friction"):
cfg.events.pop(event, None)
return cfg
+271
View File
@@ -0,0 +1,271 @@
"""Static S10 terrain courses."""
from dataclasses import dataclass
import math
import mujoco
import numpy as np
import torch
from mjlab.terrains.terrain_generator import SubTerrainCfg, TerrainOutput, TerrainGeometry, TerrainGeneratorCfg
from mjlab.managers import CurriculumTermCfg, TerminationTermCfg, RewardTermCfg
from .flat import s10_velocity_flat_env_cfg
from s10_supermimic.robot.config import (
V7_HIPX, V7_HIPY, V7_KNEE, V7_LEG_KP,
WHEEL_BODY_NAMES, WHEEL_GEOM_NAMES, KNEE_GEOM_NAMES)
STEP_HEIGHTS = (0.02, 0.04, 0.08, 0.12, 0.16, 0.20)
@dataclass
class CourseCfg(SubTerrainCfg):
kind: str = "rubble"
bidirectional: bool = False
stair_tread: float = .45
def function(self, difficulty, spec, rng):
level = min(5, max(0, round(difficulty * 5)))
h = STEP_HEIGHTS[level]
sx, sy = self.size
if self.bidirectional:
sx = 8. # Mirror the base course into a 16 m square.
body, geoms = spec.body("terrain"), []
def box(x, y, z, a, b, c, yaw=0., pitch=0., shade=.45):
# Trim the source half at its spawn plane before mirroring, avoiding
# coincident contact surfaces in the central floor/down-stair landing.
if self.bidirectional and yaw == 0. and pitch == 0.:
left, right = max(1., x-a/2), x+a/2
if right <= left:
return
x, a = (left+right)/2, right-left
# Composition Rz(yaw) Ry(pitch); z is centre, dimensions are full lengths.
cy, ss, cp, sp = math.cos(yaw/2), math.sin(yaw/2), math.cos(pitch/2), math.sin(pitch/2)
placements = [(x, (cy*cp,-ss*sp,cy*sp,ss*cp))]
if self.bidirectional:
# Reflect about x=8: yaw and pitch change sign, roll stays unchanged.
placements = [(x+7., (cy*cp,-ss*sp,cy*sp,ss*cp)),
(9.-x, (cy*cp,-ss*sp,-cy*sp,-ss*cp))]
for xx, quat in placements:
geom = body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, pos=(xx,y,z),
size=(a/2,b/2,c/2), quat=quat)
geoms.append(TerrainGeometry(geom=geom, color=(shade,shade,.9*shade,1.)))
box(sx/2,sy/2,-.1,sx,sy,.2,shade=.6)
origin = np.array([1., sy/2, 0.])
if self.kind in ("up", "down"):
n, tread = 5, self.stair_tread
if self.kind == "up":
# Overlapping foundation boxes make solid risers and a long top landing.
for k in range(n):
start = 2. + k*tread
box((start+sx)/2,sy/2,(k+1)*h/2,sx-start,sy,(k+1)*h)
else:
origin[2] = n*h
for k in range(n):
end = 2. + k*tread
height = (n-k)*h
box(end/2,sy/2,height/2,end,sy,height)
elif self.kind == "platform":
for x in (2.6,4.6):
box(x,sy/2,h/2,.9,sy,h)
elif self.kind == "slope":
# A continuous wedge made by a tilted slab with its low edge buried.
angle = math.radians(3+12*difficulty)
length = 3.
thickness = .16
center_z = length/2*math.sin(angle)-thickness/2*math.cos(angle)
box(3.5,sy/2,center_z,length,sy,thickness,pitch=-angle)
end = 3.5+length/2*math.cos(angle)
top = length*math.sin(angle)
box((end+sx)/2,sy/2,top/2,sx-end,sy,top)
elif self.kind == "rubble":
# Staggered rows cover the lane. Static overlapping rocks, never loose bodies.
# Heights of all layers together are capped by h (2--20 cm).
for ix, x in enumerate(np.arange(2.15,5.9,.48)):
for y in np.arange(.35,sy-.2,.52):
a,b = rng.uniform(.35,.65,2)
height = h*rng.uniform(.35,.7)
xx,yy = x+rng.uniform(-.12,.12),y+rng.uniform(-.12,.12)
yaw = rng.uniform(-math.pi,math.pi)
box(xx,yy,height/2,a,b,height,yaw=yaw,shade=rng.uniform(.3,.55))
if rng.random() < .25+.35*difficulty:
cap = h-height
# Deliberately overlapping cap; bounded surface height even in stacks.
pitch = rng.uniform(-.15,.15)*difficulty
half_z = (cap*math.cos(pitch)+a*.65*abs(math.sin(pitch)))/2
box(xx+.04,yy-.03,h-half_z,a*.65,b*.65,cap,
yaw=yaw+.6,pitch=pitch,shade=.55)
elif self.kind != "flat":
raise ValueError(self.kind)
if self.bidirectional:
origin[0] = 8.
return TerrainOutput(origin=origin, geometries=geoms)
def terrain_linear_tracking(env):
velocity = env.scene["robot"].data.root_link_lin_vel_b
command = env.command_manager.get_command("twist")
return torch.exp(-(velocity[:,:2]-command[:,:2]).square().sum(-1)/.25)
def terrain_yaw_tracking(env):
velocity = env.scene["robot"].data.root_link_ang_vel_b
command = env.command_manager.get_command("twist")
return torch.exp(-(velocity[:,2]-command[:,2]).square()/.5)
def course_exit(env):
offset = env.scene["robot"].data.root_link_pos_w - env.scene.env_origins
return (offset[:,0] > 5.4) | (offset[:,0] < -.65) | (offset[:,1].abs() > 2.7)
def course_curriculum(env, env_ids):
terrain = env.scene.terrain
offset = env.scene["robot"].data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
active = env.episode_length_buf[env_ids] > 0
fallen = env.termination_manager.terminated[env_ids]
success = active & ~fallen & (offset[:,0] > 5.4) & (offset[:,1].abs() < 2.7)
failure = active & (fallen | (offset[:,0] < -.65) | (offset[:,1].abs() > 2.7))
terrain.terrain_levels[env_ids] = (terrain.terrain_levels[env_ids] + success.long() - failure.long()).clamp(0,5)
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids]]
return terrain.terrain_levels.float().mean()
def bidirectional_exit(env):
offset = env.scene["robot"].data.root_link_pos_w - env.scene.env_origins
return (offset[:,0].abs() > 6.5) | (offset[:,1].abs() > 6.5)
def bidirectional_curriculum(env, env_ids):
terrain = env.scene.terrain
offset = env.scene["robot"].data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
active = env.episode_length_buf[env_ids] > 0
fallen = env.termination_manager.terminated[env_ids]
# Both course directions count. Lateral exits are neutral, not evidence
# of clearing steps; negative commands no longer count as failure.
success = active & ~fallen & (offset[:,0].abs() > 5.4) & (offset[:,1].abs() < 6.5)
failure = active & fallen
terrain.terrain_levels[env_ids] = (terrain.terrain_levels[env_ids] + success.long() - failure.long()).clamp(0,5)
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids]]
return terrain.terrain_levels.float().mean()
def s10_velocity_terrain_env_cfg(play=False, history=5, seed=0, revision=1,
actor_wheel_vel=None):
if revision not in (1, 2, 3, 4, 5, 6, 7):
raise ValueError("Terrain revision must be 1, 2, 3, 4, 5, 6 or 7")
if actor_wheel_vel is None:
actor_wheel_vel = revision <= 4 or revision >= 7
cfg = s10_velocity_flat_env_cfg(
play=play, history=history, actor_wheel_vel=actor_wheel_vel,
hipy=(V7_HIPY if revision >= 7 else None),
hipx=(V7_HIPX if revision >= 7 else None),
knee=(V7_KNEE if revision >= 7 else None),
leg_kp=(V7_LEG_KP if revision >= 7 else 80.0))
cfg.sim.nconmax = 128
cfg.sim.njmax = 512
cfg.scene.terrain.terrain_type = "generator"
cfg.scene.terrain.max_init_terrain_level = 0
cfg.scene.terrain.terrain_generator = TerrainGeneratorCfg(
seed=seed, size=(8.,8.), num_rows=6, curriculum=True, color_scheme="none",
sub_terrains={kind: CourseCfg(kind=kind, proportion=weight) for kind,weight in
(("flat",.10),("up",.25),("down",.20),("platform",.15),("slope",.10),("rubble",.20))})
cfg.events["reset_base"].params["pose_range"].update(x=(-.1,.1),y=(-.1,.1),yaw=(-.08,.08))
command = cfg.commands["twist"]
command.ranges.lin_vel_x = (.2,.45)
command.ranges.lin_vel_y = (0.,0.)
command.ranges.ang_vel_z = (-.05,.05)
command.rel_standing_envs = .1
cfg.episode_length_s = 35.
cfg.rewards["leg_posture"].weight = -.05
cfg.rewards["track_linear_velocity"].func = terrain_linear_tracking
cfg.rewards["track_linear_velocity"].params = {}
cfg.rewards["track_angular_velocity"].func = terrain_yaw_tracking
cfg.rewards["track_angular_velocity"].params = {}
cfg.rewards["upright"].params["std"] = .6
if revision >= 2:
from s10_supermimic.core.terrain_posture import env_posture_cost
cfg.actions["joint_pos"].scale = {".*_hipx_joint": .25, ".*_hipy_joint": .6, ".*_knee_joint": .8}
cfg.rewards["leg_posture"].func = env_posture_cost
cfg.rewards["leg_posture"].weight = -.4
cfg.curriculum = {} if play else {"course": CurriculumTermCfg(func=course_curriculum)}
cfg.terminations["course_exit"] = TerminationTermCfg(func=course_exit, time_out=True)
if revision >= 3:
command.ranges.lin_vel_x = (-2., 2.)
command.ranges.lin_vel_y = (-1., 1.)
command.ranges.ang_vel_z = (-3., 3.)
command.rel_forward_envs = 0.
command.rel_world_envs = 0.
cfg.actions["wheel_vel"].scale = 35.
gen = cfg.scene.terrain.terrain_generator
gen.size = (16., 16.)
for course in gen.sub_terrains.values():
course.bidirectional = True
cfg.curriculum = {} if play else {"course": CurriculumTermCfg(func=bidirectional_curriculum)}
cfg.terminations["course_exit"] = TerminationTermCfg(func=bidirectional_exit, time_out=True)
if revision >= 4:
from mjlab.sensor import ContactMatch, ContactSensorCfg
cfg.scene.sensors = tuple(cfg.scene.sensors or ()) + (ContactSensorCfg(
name="wheel_ground_contact", primary=ContactMatch(mode="body",
pattern=tuple(WHEEL_BODY_NAMES), entity="robot"),
secondary=ContactMatch(mode="body", pattern="terrain"),
fields=("found","force"), reduce="netforce", num_slots=1),)
from s10_supermimic.core.rolling_rewards import env_rolling_cost
for course in cfg.scene.terrain.terrain_generator.sub_terrains.values():
course.stair_tread = .30
for name, weight in (("unnecessary_wheel_lift", -.2), ("rolling_leg_motion", -.12)):
cfg.rewards[name] = RewardTermCfg(func=env_rolling_cost, weight=weight, params={"term":name})
if not play:
from s10_supermimic.core.mixed_commands import MixedVelocityCommandCfg
cfg.commands["twist"] = MixedVelocityCommandCfg(**vars(command))
if revision >= 5:
from s10_supermimic.core.obstacle_rewards import env_obstacle_progress
from s10_supermimic.core.stair_curriculum import curriculum
cfg.rewards['obstacle_progress'] = RewardTermCfg(func=env_obstacle_progress, weight=.5)
if not play:
cfg.curriculum = {'course': CurriculumTermCfg(func=curriculum)}
cfg.commands['twist'].hold_stair_command = True
for kind, weight in zip(('flat','up','down','platform','slope','rubble'), (.1,.35,.15,.1,.1,.2)):
cfg.scene.terrain.terrain_generator.sub_terrains[kind].proportion = weight
if revision == 6:
from s10_supermimic.core.wheel_terrain_rewards import env_wheel_term
for name in ('leg_posture', 'unnecessary_wheel_lift', 'rolling_leg_motion'):
cfg.rewards[name].func = env_wheel_term
cfg.rewards[name].params = {'term': name}
cfg.rewards['wheel_obstacle_clearance'] = RewardTermCfg(
func=env_wheel_term, weight=-.25, params={'term':'wheel_obstacle_clearance'})
if revision >= 7:
from s10_supermimic.core.collision import illegal_nonwheel_contact
from s10_supermimic.core.mixed_commands import MixedVelocityCommandCfg
from s10_supermimic.core.wheel_intrusion import (
env_pair_body_intrusion, env_wheel_body_intrusion)
# V7 = V5 + non-wheel collisions; experimental V6 shaping is independent.
cfg.scene.sensors += (ContactSensorCfg(
name='nonwheel_ground_contact',
primary=ContactMatch(mode='geom', pattern='.*_collision', entity='robot',
exclude=WHEEL_GEOM_NAMES + KNEE_GEOM_NAMES),
secondary=ContactMatch(mode='body', pattern='terrain'),
fields=('found','force'), reduce='netforce', num_slots=1,
history_length=cfg.decimation),)
cfg.terminations['illegal_nonwheel_contact'] = TerminationTermCfg(
func=illegal_nonwheel_contact, params={'force_threshold':50.})
# Safe low-speed all-terrain envelope. Fixed for training and play so the
# validation videos use exactly the command distribution seen in training.
command = cfg.commands['twist']
command.ranges.lin_vel_x = (-1.5, 1.5)
command.ranges.lin_vel_y = (-.8, .8)
command.ranges.ang_vel_z = (-1.2, 1.2)
if not isinstance(command, MixedVelocityCommandCfg):
command = MixedVelocityCommandCfg(**vars(command))
command.deadzone = .05
command.hold_stair_command = not play
cfg.commands['twist'] = command
# V7 initialization: small x/y/yaw only, no random roll/pitch.
# Always active, independent of terrain relief: keep front/rear wheels out
# of the opposite two-thirds of the body projection.
cfg.rewards['wheel_body_intrusion'] = RewardTermCfg(
func=env_wheel_body_intrusion, weight=-1.0)
# Wider 5/6 restriction: both same-side legs may not sit under the body
# together for more than 0.5 s. One leg passing through is free.
cfg.rewards['pair_body_intrusion'] = RewardTermCfg(
func=env_pair_body_intrusion, weight=-1.0)
# Stairs need deep flexion. Keep hipx constrained (2x in posture_cost) but
# let hipy/knee deviate almost freely once local step relief is large.
cfg.rewards['leg_posture'].params = {
**cfg.rewards['leg_posture'].params, 'relief_floor': .02}
return cfg