150 lines
6.0 KiB
Python
150 lines
6.0 KiB
Python
"""S10 actor history and actor/critic networks, independent of the quadruped learner.
|
|
|
|
Two Actor observation schemas coexist:
|
|
|
|
* v4 / v4-upward / new v7: 53 values/frame, including measured wheel speed;
|
|
* legacy v7: 49 values/frame, no measured wheel speed (the last action still
|
|
contains the four wheel commands).
|
|
|
|
The learner, replay and deployment exporter select the schema explicitly. Old
|
|
checkpoints are recognized from their first-layer width.
|
|
"""
|
|
|
|
import torch
|
|
from torch import nn
|
|
from .world_model import RunningNorm, ANG_VEL, GRAVITY, LEG_POS, LEG_VEL, WHEEL_VEL
|
|
|
|
LEGACY_TERMS = (3, 3, 12, 12, 4, 16, 3)
|
|
TERMS = (3, 3, 12, 12, 16, 3)
|
|
LEGACY_FRAME_DIM = sum(LEGACY_TERMS)
|
|
FRAME_DIM = sum(TERMS)
|
|
CRITIC_DIM = 38 + 3 + 7 + 16
|
|
POLICY_FORMAT = "s10-flat-policy-v1"
|
|
|
|
|
|
def frame(state, last_env_action, command, bias, noise=False, wheel_vel=True):
|
|
"""Exactly the term order/scaling in env.flat (noise before scaling)."""
|
|
values = [state[:, ANG_VEL], state[:, GRAVITY], state[:, LEG_POS] + bias,
|
|
state[:, LEG_VEL]]
|
|
widths = [0.2, 0.05, 0.01, 1.5]
|
|
if wheel_vel:
|
|
values.append(state[:, WHEEL_VEL])
|
|
widths.append(1.5)
|
|
values.extend((last_env_action, command))
|
|
widths.extend((0., 0.))
|
|
if noise:
|
|
values = [v + (torch.rand_like(v) * 2 - 1) * w if w else v
|
|
for v, w in zip(values, widths)]
|
|
# Leg velocity is always index 3; wheel velocity, when present, is index 4.
|
|
velocity_indices = (3, 4) if wheel_vel else (3,)
|
|
return torch.cat([v * (0.05 if i in velocity_indices else 1.)
|
|
for i, v in enumerate(values)], -1)
|
|
|
|
|
|
def push_history(history, latest, length, terms):
|
|
"""Term-wise histories, oldest to newest, matching mjlab flattening."""
|
|
result, hidx, fidx = [], 0, 0
|
|
for dim in terms:
|
|
result.append(torch.cat((history[:, hidx + dim:hidx + dim * length],
|
|
latest[:, fidx:fidx + dim]), -1))
|
|
hidx += dim * length
|
|
fidx += dim
|
|
return torch.cat(result, -1)
|
|
|
|
|
|
def critic_input(state, command, context, previous):
|
|
return torch.cat((state, command, context, previous), -1)
|
|
|
|
|
|
class Network(nn.Module):
|
|
def __init__(self, inputs, outputs, hidden):
|
|
super().__init__()
|
|
self.norm = RunningNorm(inputs)
|
|
layers = []
|
|
for n in hidden:
|
|
layers.extend((nn.Linear(inputs, n), nn.ELU()))
|
|
inputs = n
|
|
layers.append(nn.Linear(inputs, outputs))
|
|
self.net = nn.Sequential(*layers)
|
|
|
|
def forward(self, x):
|
|
return self.net(self.norm(x))
|
|
|
|
|
|
class Actor(Network):
|
|
def __init__(self, history=5, hidden=(256, 128, 64), actor_wheel_vel=True):
|
|
self.actor_wheel_vel = bool(actor_wheel_vel)
|
|
self.terms = LEGACY_TERMS if self.actor_wheel_vel else TERMS
|
|
self.frame_dim = sum(self.terms)
|
|
super().__init__(self.frame_dim * history, 16, hidden)
|
|
self.history, self.hidden = history, tuple(hidden)
|
|
nn.init.uniform_(self.net[-1].weight, -0.01, 0.01)
|
|
nn.init.zeros_(self.net[-1].bias)
|
|
|
|
def forward(self, obs):
|
|
# Same bounded action distribution during collection and imagination.
|
|
return super().forward(obs).tanh()
|
|
|
|
|
|
def checkpoint_wheel_vel(checkpoint):
|
|
"""Infer legacy observation schema from weights; validate explicit metadata."""
|
|
flag = checkpoint.get("actor_wheel_vel")
|
|
if "actor" in checkpoint:
|
|
dim = checkpoint["actor"]["net.0.weight"].shape[1]
|
|
history = checkpoint["history"]
|
|
if dim not in (FRAME_DIM * history, LEGACY_FRAME_DIM * history):
|
|
raise ValueError("Unsupported actor observation dimension")
|
|
inferred = dim == LEGACY_FRAME_DIM * history
|
|
if flag is not None and bool(flag) != inferred:
|
|
raise ValueError("Actor observation metadata disagrees with weights")
|
|
return inferred
|
|
return bool(flag) if flag is not None else False
|
|
|
|
|
|
def load_actor(checkpoint, device="cpu"):
|
|
if checkpoint.get("format") not in (POLICY_FORMAT, "s10-terrain-policy-v1"):
|
|
raise ValueError("Expected a S10 policy checkpoint, not other-robot or dynamics-only weights")
|
|
actor = Actor(checkpoint["history"], checkpoint["actor_hidden"],
|
|
checkpoint_wheel_vel(checkpoint)).to(device)
|
|
actor.load_state_dict(checkpoint["actor"])
|
|
return actor.eval()
|
|
|
|
|
|
def evaluation_env_cfg(checkpoint, level=0, kind=None):
|
|
"""Select the matching task; terrain remains privileged at deployment."""
|
|
from s10_supermimic.tasks import get_task
|
|
actor_wheel_vel = checkpoint_wheel_vel(checkpoint)
|
|
if checkpoint.get("format") == "s10-terrain-policy-v1":
|
|
checkpoint_task = checkpoint.get("config", {}).get("task")
|
|
task = get_task(checkpoint_task) if checkpoint_task else None
|
|
if task is not None:
|
|
if task.terrain_revision != checkpoint["config"].get("terrain_revision", 1):
|
|
raise ValueError("Checkpoint task revision differs from its config")
|
|
if task.actor.wheel_vel != actor_wheel_vel:
|
|
raise ValueError("Checkpoint task Actor schema differs from its weights")
|
|
cfg = task.build_env_cfg(play=True, history=checkpoint["history"],
|
|
seed=checkpoint["config"]["seed"])
|
|
else:
|
|
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
|
cfg = s10_velocity_terrain_env_cfg(
|
|
play=True, history=checkpoint["history"],
|
|
seed=checkpoint["config"]["seed"],
|
|
revision=checkpoint["config"].get("terrain_revision", 1),
|
|
actor_wheel_vel=actor_wheel_vel)
|
|
# A single staircase can overlap the robot's pre-reset origin at compile
|
|
# time. Reserve contacts for that transient pose as well as normal rolling.
|
|
cfg.sim.nconmax = max(cfg.sim.nconmax or 0, 512)
|
|
cfg.sim.njmax = max(cfg.sim.njmax or 0, 2048)
|
|
if level not in range(6):
|
|
raise ValueError("Terrain level must be 0..5")
|
|
gen = cfg.scene.terrain.terrain_generator
|
|
# Evaluation fixes difficulty; one row, same seeded generator.
|
|
gen.num_rows = 1
|
|
gen.difficulty_range = (level/5,level/5)
|
|
if kind is not None:
|
|
gen.sub_terrains = {kind: gen.sub_terrains[kind]}
|
|
return cfg
|
|
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
|
return s10_velocity_flat_env_cfg(play=True, history=checkpoint["history"],
|
|
actor_wheel_vel=actor_wheel_vel)
|