76 lines
3.5 KiB
Python
76 lines
3.5 KiB
Python
"""V5 episode outcomes and performance-gated velocity curriculum.
|
|
|
|
All buffers are training-only metadata, never actor observations.
|
|
"""
|
|
import torch
|
|
|
|
KINDS = ('flat', 'up', 'down', 'platform', 'slope', 'rubble')
|
|
MIN_ITERATIONS = (0, 3000, 8000, 15000, 22000)
|
|
|
|
|
|
def buffers(env):
|
|
if not hasattr(env, '_stair_stats'):
|
|
env._stair_stats = torch.zeros(6, 6, 4, device=env.device) # trials, pass, fall, stuck
|
|
env._stair_stall_time = torch.zeros(env.num_envs, device=env.device)
|
|
env._stair_stuck = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
|
return env._stair_stats
|
|
|
|
|
|
def stair_pass(offset, gravity, level, up):
|
|
"""Require body/rear clearance of the last riser, correct landing height and stability."""
|
|
heights = offset.new_tensor([.02, .04, .08, .12, .16, .20])[level]
|
|
expected = torch.where(up, 5 * heights, -5 * heights)
|
|
return ((offset[:, 0].abs() > 2.8) & (offset[:, 1].abs() < 2.5)
|
|
& (offset[:, 2] > expected + .20) & (offset[:, 2] < expected + .60)
|
|
& (gravity[:, :2].norm(dim=-1) < .4))
|
|
|
|
|
|
def curriculum(env, env_ids):
|
|
stats = buffers(env)
|
|
terrain = env.scene.terrain
|
|
d = env.scene['robot'].data
|
|
offset = d.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
|
levels = terrain.terrain_levels[env_ids].clone()
|
|
types = terrain.terrain_types[env_ids]
|
|
# Six fixed columns, generated in KINDS order.
|
|
stairs = (types == 1) | (types == 2)
|
|
active = env.episode_length_buf[env_ids] > 0
|
|
fallen = env.termination_manager.terminated[env_ids]
|
|
success = (offset[:, 0].abs() > 5.4) & (offset[:, 1].abs() < 6.5)
|
|
top = stair_pass(offset, d.projected_gravity_b[env_ids], levels, types == 1)
|
|
# Landing contacts rule out simply jumping through the clearance test.
|
|
contact = env.scene['wheel_ground_contact'].data.found[env_ids].reshape(len(env_ids), -1).gt(0).sum(-1) >= 2
|
|
success = torch.where(stairs, top & contact, success) & active & ~fallen
|
|
command = env.command_manager.get_command('twist')[env_ids]
|
|
traversal = (command[:, 0].abs() > .1) & (command[:, 1:].abs().amax(-1) < .05)
|
|
trials = (~stairs | traversal) & active & ((env.episode_length_buf[env_ids] * env.step_dt >= 5.) | fallen)
|
|
values = torch.stack((trials, success & trials, fallen & trials,
|
|
env._stair_stuck[env_ids] & trials), -1).float()
|
|
stats.view(-1, 4).index_add_(0, types * 6 + levels, values)
|
|
terrain.terrain_levels[env_ids] = (levels + success.long() - (active & fallen).long()).clamp(0, 5)
|
|
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], types]
|
|
env._stair_stall_time[env_ids] = 0
|
|
env._stair_stuck[env_ids] = False
|
|
return terrain.terrain_levels.float().mean()
|
|
|
|
|
|
def next_stage(stage, iteration, stats, velocity_mae):
|
|
"""Minimum age AND sufficient actual stair passes, never time alone."""
|
|
if stage >= 4 or iteration < MIN_ITERATIONS[stage + 1]:
|
|
return stage
|
|
# Stage increases require success at 4, 8, 12, then 16 cm respectively.
|
|
trials, passed = stats[1, stage + 1, :2]
|
|
accurate = max(velocity_mae[:2]) < .18 and velocity_mae[2] < .30
|
|
return stage + int(trials >= 64 and passed / max(trials, 1) >= .60 and accurate)
|
|
|
|
|
|
def summary(env):
|
|
stats = buffers(env).detach().cpu().tolist()
|
|
levels, types = env.scene.terrain.terrain_levels, env.scene.terrain.terrain_types
|
|
result = {}
|
|
for i, name in enumerate(KINDS):
|
|
selected = levels[types == i]
|
|
result[name] = dict(level_counts=torch.bincount(selected, minlength=6).cpu().tolist(),
|
|
episodes_by_level=stats[i])
|
|
return result
|