Initial commit

2026-11-02 17:10:39 +08:00
commit 102fdbaf47
79 changed files with 56509 additions and 0 deletions
@@ -0,0 +1,34 @@
+"""Bounded directional motion near uneven ground, shared by real and imagined steps."""
+import torch
+from .world_model import LIN_VEL
+
+
+def obstacle_gate(scan):
+  local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
+  return ((local.amax(-1) - local.amin(-1) - .02) / .08).clamp(0., 1.)
+
+
+def obstacle_progress(state, command, scan):
+  speed = command[:, :2].norm(dim=-1)
+  along = (state[:, LIN_VEL][:, :2] * command[:, :2]).sum(-1) / speed.clamp_min(.1)
+  # Signed and capped: no benefit for wheel spin, pure yaw, lifting or overspeed.
+  return obstacle_gate(scan) * (speed > .1) * (along / speed.clamp_min(.2)).clamp(-1., 1.)
+
+
+def env_obstacle_progress(env):
+  from .terrain import TerrainField, TerrainBridge
+  from .stair_curriculum import buffers
+  bridge = getattr(env, '_s10_posture_bridge', None)
+  if bridge is None:
+    bridge = TerrainBridge(env, TerrainField.from_env(env))
+    env._s10_posture_bridge = bridge
+  state = bridge.state()
+  scan = bridge.field.scan(state, bridge.pose())
+  command = env.command_manager.get_command('twist')
+  buffers(env)
+  stalled = ((command[:, :2].norm(dim=-1) > .1)
+             & (state[:, LIN_VEL][:, :2].norm(dim=-1) < .05)
+             & (obstacle_gate(scan) > .5))
+  env._stair_stall_time = torch.where(stalled, env._stair_stall_time + env.step_dt, 0.)
+  env._stair_stuck |= env._stair_stall_time >= 3.
+  return obstacle_progress(state, command, scan)