Add Rolling Start talent policy, timeline and training recipe

b791e4c verified 4 days ago

11 kB

"""Episode-return policy search for the head-to-wheel recovery on the Arena's roller rig."""

import functools
import hashlib
import json
import multiprocessing as mp
import os
from pathlib import Path

import numpy as np
from arena import runner
from arena.course import compile_stage
from arena.manifest import load_event
from arena.policy import load_policy
from arena.timeline import Start, Timeline

OUT = Path(__file__).resolve().parent
RIG = "rollers"
EVENT = load_event("got-a-talent")
START = Start(x=0, y=-1.45, yaw_deg=90)
BASE_MODELS = {
    "skate": (
        "JackDH/microduck-roller-sprint-20261110-best",
        "43e9b57825fbfe08cf4b510368b643ff9ae47d93",
        "policy.onnx",
        "0b724f5cebe7f50815f6dd9631009e47c4b0a607f81b22d0f4048ffd6d5a77d6",
    ),
    "roll": (
        "pollen-robotics/microduck-policies",
        "1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
        "roulade.onnx",
        "3d60da08fc13f29c1b57f41977aa898132c0d60042100149d8e775affcbca32b",
    ),
    "rise": (
        "pollen-robotics/microduck-policies",
        "1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
        "velstand.onnx",
        "1c659be55da94bc5753b707de5c6a3e7c49931e05ca3b6991615cef1a8ba9a45",
    ),
}


@functools.lru_cache(maxsize=3)
def model_file(name):
    repo, revision, filename, digest = BASE_MODELS[name]
    path = None
    if project := os.environ.get("MICRODUCK_PROJECT"):
        rel = (
            "challenges/publish-roller-best-20261110/policy.onnx"
            if name == "skate"
            else f"policies/official/{filename}"
        )
        path = Path(project) / rel
    if path is None or not path.is_file():
        from huggingface_hub import hf_hub_download

        path = Path(hf_hub_download(repo, filename, revision=revision))
    if hashlib.sha256(path.read_bytes()).hexdigest() != digest:
        raise ValueError(f"Unexpected contents for {name} policy")
    return path


# The timeline gives the coarse cues. The network reads gravity and angular velocity to adjust its leg action.
TIMES = [0, 0.2, 1.37, 1.55, 1.87, 2.33, 2.85, 3.15]
MODES = [0, 0, 1, 2, 3, 4, 5, 0]
SPEEDS = [0, 0.32, 0, 0, 0, 0.54, 0.54, 0.65]
DURATION = 5.16


def timeline_dict():
    frames = [
        {"t": t, "twist": [v, 0, 0], "body": [0, 0, 0, 0, 0, m]}
        for t, v, m in zip(TIMES, SPEEDS, MODES)
    ]
    return {
        "timeline_version": 2,
        "duration_s": DURATION,
        "ducks": [{"policy": "self", "start": START.to_dict(), "keyframes": frames}],
    }


TIMELINE = Timeline(
    DURATION,
    tuple(TIMES),
    tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
)


@functools.lru_cache(maxsize=4)
def actor(name):
    p = load_policy(model_file(name))
    return p.session, p._input_name, p._output_name


def actions(name, x):
    sess, i, o = actor(name)
    return sess.run([o], {i: x.reshape(1, 61)})[0][0]


# Symmetric sagittal legs, neck/head, and a lateral hip balance response.
MAP = np.zeros((6, 14), np.float32)
for n, (l, r) in enumerate([(2, 11), (3, 12), (4, 13)]):
    MAP[n, l] = 1
    MAP[n, r] = -1
MAP[3, 5] = 1
MAP[4, 6] = 1
MAP[5, [1, 10]] = [1, -1]
SHAPE = (4, 6)
SIZE = 32


def configure(speed=0.32, entry=1.37):
    global TIMES, SPEEDS, DURATION, TIMELINE
    offset = entry - 1.37
    TIMES = [0, 0.2] + [
        round(t + offset, 2) for t in (1.37, 1.55, 1.87, 2.33, 2.85, 3.15)
    ]
    SPEEDS = [0, speed, 0, 0, 0, 0.54, 0.54, 0.65]
    DURATION = round(5.16 + offset, 2)
    TIMELINE = Timeline(
        DURATION,
        tuple(TIMES),
        tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
    )


class Landing:
    def __init__(self, theta):
        self.p = np.asarray(theta, np.float32)
        self.steps = self.p[:24].reshape(SHAPE)

    def act(self, obs):
        mode = int(obs[-1] + 0.5)
        x = obs.copy()
        x[49:] = 0
        if mode == 0:
            return actions("skate", x)
        if mode <= 3:
            x[34:48] /= np.float32(1.12)
            x[48] = 0
            base = actions("roll", x) * np.float32(1.12)
        else:
            base = actions("rise", x)
        if mode == 1:
            return base
        residual = self.steps[mode - 2].copy()
        # The last eight parameters adapt the joints to pitch, roll and their velocities.
        pitch_vel = np.clip(obs[1] * np.float32(0.14), -1, 1)
        residual[:3] += self.p[24:27] * obs[3] + self.p[27:30] * pitch_vel
        residual[5] += self.p[30] * obs[4] + self.p[31] * np.clip(
            obs[0] * np.float32(0.15), -1, 1
        )
        return base + np.tanh(residual) @ MAP


# Compile each stage once per process. Each new game resets friction/delay and carries the original physics.
@functools.lru_cache(maxsize=1)
def stage():
    return compile_stage(EVENT, [(RIG, START)])


def run(policy, seed):
    model = stage()
    sim, settle, _ = runner._start(
        EVENT,
        model,
        [policy],
        [START.yaw_deg],
        int(seed),
        [np.asarray(TIMELINE.command_at(0))],
    )
    trace = np.empty((round(DURATION * 50), 21), np.float32)
    for i in range(len(trace)):
        sim.step([np.asarray(TIMELINE.command_at(i * 0.02))])
        sim.record(trace[i : i + 1])
        if runner._reset_by_mujoco(sim.data):
            raise RuntimeError("Diverged")
    return trace, settle


def metrics(tr):
    w, x, y, z = tr[:, 3:7].T
    tilt = np.degrees(np.arccos(np.clip(1 - 2 * (x * x + y * y), -1, 1)))
    pitch = np.degrees(np.unwrap(np.arctan2(2 * (w * y - x * z), 2 * (x * y + w * z))))
    start = int(np.ceil(TIMES[2] * 50 - 1e-8)) - 1
    angle = pitch - np.median(pitch[start - 8 : start])
    theta = float(np.median(angle[-25:]))
    h = tr[:, 2]
    pos = tr[:, 1]
    upright = (tilt < 35) & (h > 0.085)
    inv = float(tilt[start:].max())
    glide = float(pos[-1] - pos[-51])
    inside = bool((abs(tr[:, 0]) < 1.28).all() and (abs(pos) < 1.87).all())
    success = bool(
        310 < theta < 405
        and inv > 150
        and np.all(upright[-45:])
        and glide > 0.09
        and inside
        and pos[start] - pos[start - 25] > 0.035
    )
    # Reward partial catches continuously and reserve the completion bonus for a genuine roll and rollout.
    turn = np.exp(-(((angle - 360) / 82) ** 2))
    orient = np.clip((1 + 1 - 2 * (x * x + y * y)) / 2, 0, 1)
    height = np.clip((h - 0.043) / 0.073, 0, 1)
    alignment = np.clip(2 * (x * y + w * z), 0, 1)
    quality = turn * orient * height * (0.45 + 0.55 * alignment)
    reward = (
        0.9 * np.clip((inv - 100) / 70, 0, 1)
        + 1.3 * quality[start + 50 : start + 100].mean()
        + 3.1 * quality[start + 100 :].mean()
        + 1.4 * np.clip(glide / 0.36, -0.5, 1) * quality[-50:].mean()
        + 2.5 * success
    )
    good = np.flatnonzero(~upright[start:])
    recovery = round((int(good[-1]) + 1) * 0.02, 2) if len(good) and success else None
    return {
        "return": round(float(reward), 6),
        "success": success,
        "turn_deg": round(theta, 1),
        "max_tilt": round(inv, 1),
        "entry_mps": round(float((pos[start] - pos[start - 25]) / 0.5), 3),
        "entry_distance_m": round(float(pos[start] - pos[0]), 3),
        "final_m": round(glide, 3),
        "recovery_s": recovery,
        "inside": inside,
    }


def test_job(args):
    theta, seed = args
    tr, _ = run(Landing(theta), seed)
    return metrics(tr)


def calc(pool, models, seeds):
    res = list(
        pool.imap(test_job, [(v, s) for v in models for s in seeds], chunksize=2)
    )
    return [res[i : i + len(seeds)] for i in range(0, len(res), len(seeds))]


def save(path, obj):
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(obj, indent=2) + "\n")
    tmp.replace(path)


def train(
    steps,
    pop,
    batch,
    warm=None,
    run_dir=OUT,
    random_seed=11062026,
    validation=None,
    search_scale=0.62,
    floor=0.075,
    workers=16,
):
    rng = np.random.default_rng(random_seed)
    mean = np.zeros(SIZE)
    sig = np.full(SIZE, 0.33)
    sig[24:] = 0.22
    best = []
    hist = []
    if warm is not None:
        mean = np.asarray(warm)
        sig *= search_scale
    with mp.Pool(
        workers, initializer=configure, initargs=(SPEEDS[1], TIMES[2])
    ) as pool:
        validation = list(range(31, 43)) if validation is None else list(validation)
        baseline = calc(pool, [mean], validation)[0]
        print(
            "Baseline validation:",
            sum(r["success"] for r in baseline),
            "/",
            len(baseline),
            flush=True,
        )
        for gen in range(steps):
            first = rng.normal(size=(pop // 2, SIZE))
            eps = np.concatenate((first, -first))
            candidates = np.clip(mean + sig * eps, -2.8, 2.8)
            candidates[0] = mean
            seeds = list(
                map(int, rng.choice(np.arange(80, 600), size=batch, replace=False))
            )
            results = calc(pool, candidates, seeds)
            returns = np.array([np.mean([a["return"] for a in rr]) for rr in results])
            elite = candidates[np.argsort(returns)[-max(5, pop // 6) :]]
            mean = 0.27 * mean + 0.73 * np.mean(elite, axis=0)
            sig = np.maximum(floor, 0.62 * sig + 0.38 * np.std(elite, axis=0))
            leaders = [mean, candidates[np.argmax(returns)]]
            if (gen + 1) % 2 == 0 or gen == 0:
                trials = calc(pool, leaders, [*validation, 1])
                for model, rr in zip(leaders, trials):
                    item = {
                        "gen": gen + 1,
                        "theta": model.tolist(),
                        "validation_success": sum(q["success"] for q in rr[:-1]),
                        "official_pass": rr[-1]["success"],
                        "validation_return": float(
                            np.mean([q["return"] for q in rr[:-1]])
                        ),
                    }
                    best.append(item)
                best.sort(
                    key=lambda o: (o["validation_success"], o["validation_return"]),
                    reverse=True,
                )
                best = best[:5] + [o for o in best[5:] if o["official_pass"]][:3]
                save(run_dir / "checkpoints.json", best)
            hist.append(
                {
                    "generation": gen + 1,
                    "training_seeds": seeds,
                    "mean_return": float(returns.mean()),
                    "best_training_return": float(returns.max()),
                }
            )
            save(run_dir / "learning.json", hist)
            print(
                "Generation",
                gen + 1,
                "best validation",
                best[0]["validation_success"],
                "/",
                len(validation),
                "training return",
                round(returns.max(), 2),
                flush=True,
            )

	"""Episode-return policy search for the head-to-wheel recovery on the Arena's roller rig."""

	import functools
	import hashlib
	import json
	import multiprocessing as mp
	import os
	from pathlib import Path

	import numpy as np
	from arena import runner
	from arena.course import compile_stage
	from arena.manifest import load_event
	from arena.policy import load_policy
	from arena.timeline import Start, Timeline

	OUT = Path(__file__).resolve().parent
	RIG = "rollers"
	EVENT = load_event("got-a-talent")
	START = Start(x=0, y=-1.45, yaw_deg=90)
	BASE_MODELS = {
	"skate": (
	"JackDH/microduck-roller-sprint-20261110-best",
	"43e9b57825fbfe08cf4b510368b643ff9ae47d93",
	"policy.onnx",
	"0b724f5cebe7f50815f6dd9631009e47c4b0a607f81b22d0f4048ffd6d5a77d6",
	),
	"roll": (
	"pollen-robotics/microduck-policies",
	"1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
	"roulade.onnx",
	"3d60da08fc13f29c1b57f41977aa898132c0d60042100149d8e775affcbca32b",
	),
	"rise": (
	"pollen-robotics/microduck-policies",
	"1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
	"velstand.onnx",
	"1c659be55da94bc5753b707de5c6a3e7c49931e05ca3b6991615cef1a8ba9a45",
	),
	}


	@functools.lru_cache(maxsize=3)
	def model_file(name):
	repo, revision, filename, digest = BASE_MODELS[name]
	path = None
	if project := os.environ.get("MICRODUCK_PROJECT"):
	rel = (
	"challenges/publish-roller-best-20261110/policy.onnx"
	if name == "skate"
	else f"policies/official/{filename}"
	)
	path = Path(project) / rel
	if path is None or not path.is_file():
	from huggingface_hub import hf_hub_download

	path = Path(hf_hub_download(repo, filename, revision=revision))
	if hashlib.sha256(path.read_bytes()).hexdigest() != digest:
	raise ValueError(f"Unexpected contents for {name} policy")
	return path


	# The timeline gives the coarse cues. The network reads gravity and angular velocity to adjust its leg action.
	TIMES = [0, 0.2, 1.37, 1.55, 1.87, 2.33, 2.85, 3.15]
	MODES = [0, 0, 1, 2, 3, 4, 5, 0]
	SPEEDS = [0, 0.32, 0, 0, 0, 0.54, 0.54, 0.65]
	DURATION = 5.16


	def timeline_dict():
	frames = [
	{"t": t, "twist": [v, 0, 0], "body": [0, 0, 0, 0, 0, m]}
	for t, v, m in zip(TIMES, SPEEDS, MODES)
	]
	return {
	"timeline_version": 2,
	"duration_s": DURATION,
	"ducks": [{"policy": "self", "start": START.to_dict(), "keyframes": frames}],
	}


	TIMELINE = Timeline(
	DURATION,
	tuple(TIMES),
	tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
	)


	@functools.lru_cache(maxsize=4)
	def actor(name):
	p = load_policy(model_file(name))
	return p.session, p._input_name, p._output_name


	def actions(name, x):
	sess, i, o = actor(name)
	return sess.run([o], {i: x.reshape(1, 61)})[0][0]


	# Symmetric sagittal legs, neck/head, and a lateral hip balance response.
	MAP = np.zeros((6, 14), np.float32)
	for n, (l, r) in enumerate([(2, 11), (3, 12), (4, 13)]):
	MAP[n, l] = 1
	MAP[n, r] = -1
	MAP[3, 5] = 1
	MAP[4, 6] = 1
	MAP[5, [1, 10]] = [1, -1]
	SHAPE = (4, 6)
	SIZE = 32


	def configure(speed=0.32, entry=1.37):
	global TIMES, SPEEDS, DURATION, TIMELINE
	offset = entry - 1.37
	TIMES = [0, 0.2] + [
	round(t + offset, 2) for t in (1.37, 1.55, 1.87, 2.33, 2.85, 3.15)
	]
	SPEEDS = [0, speed, 0, 0, 0, 0.54, 0.54, 0.65]
	DURATION = round(5.16 + offset, 2)
	TIMELINE = Timeline(
	DURATION,
	tuple(TIMES),
	tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
	)


	class Landing:
	def __init__(self, theta):
	self.p = np.asarray(theta, np.float32)
	self.steps = self.p[:24].reshape(SHAPE)

	def act(self, obs):
	mode = int(obs[-1] + 0.5)
	x = obs.copy()
	x[49:] = 0
	if mode == 0:
	return actions("skate", x)
	if mode <= 3:
	x[34:48] /= np.float32(1.12)
	x[48] = 0
	base = actions("roll", x) * np.float32(1.12)
	else:
	base = actions("rise", x)
	if mode == 1:
	return base
	residual = self.steps[mode - 2].copy()
	# The last eight parameters adapt the joints to pitch, roll and their velocities.
	pitch_vel = np.clip(obs[1] * np.float32(0.14), -1, 1)
	residual[:3] += self.p[24:27] * obs[3] + self.p[27:30] * pitch_vel
	residual[5] += self.p[30] * obs[4] + self.p[31] * np.clip(
	obs[0] * np.float32(0.15), -1, 1
	)
	return base + np.tanh(residual) @ MAP


	# Compile each stage once per process. Each new game resets friction/delay and carries the original physics.
	@functools.lru_cache(maxsize=1)
	def stage():
	return compile_stage(EVENT, [(RIG, START)])


	def run(policy, seed):
	model = stage()
	sim, settle, _ = runner._start(
	EVENT,
	model,
	[policy],
	[START.yaw_deg],
	int(seed),
	[np.asarray(TIMELINE.command_at(0))],
	)
	trace = np.empty((round(DURATION * 50), 21), np.float32)
	for i in range(len(trace)):
	sim.step([np.asarray(TIMELINE.command_at(i * 0.02))])
	sim.record(trace[i : i + 1])
	if runner._reset_by_mujoco(sim.data):
	raise RuntimeError("Diverged")
	return trace, settle


	def metrics(tr):
	w, x, y, z = tr[:, 3:7].T
	tilt = np.degrees(np.arccos(np.clip(1 - 2 * (x * x + y * y), -1, 1)))
	pitch = np.degrees(np.unwrap(np.arctan2(2 * (w * y - x * z), 2 * (x * y + w * z))))
	start = int(np.ceil(TIMES[2] * 50 - 1e-8)) - 1
	angle = pitch - np.median(pitch[start - 8 : start])
	theta = float(np.median(angle[-25:]))
	h = tr[:, 2]
	pos = tr[:, 1]
	upright = (tilt < 35) & (h > 0.085)
	inv = float(tilt[start:].max())
	glide = float(pos[-1] - pos[-51])
	inside = bool((abs(tr[:, 0]) < 1.28).all() and (abs(pos) < 1.87).all())
	success = bool(
	310 < theta < 405
	and inv > 150
	and np.all(upright[-45:])
	and glide > 0.09
	and inside
	and pos[start] - pos[start - 25] > 0.035
	)
	# Reward partial catches continuously and reserve the completion bonus for a genuine roll and rollout.
	turn = np.exp(-(((angle - 360) / 82) ** 2))
	orient = np.clip((1 + 1 - 2 * (x * x + y * y)) / 2, 0, 1)
	height = np.clip((h - 0.043) / 0.073, 0, 1)
	alignment = np.clip(2 * (x * y + w * z), 0, 1)
	quality = turn * orient * height * (0.45 + 0.55 * alignment)
	reward = (
	0.9 * np.clip((inv - 100) / 70, 0, 1)
	+ 1.3 * quality[start + 50 : start + 100].mean()
	+ 3.1 * quality[start + 100 :].mean()
	+ 1.4 * np.clip(glide / 0.36, -0.5, 1) * quality[-50:].mean()
	+ 2.5 * success
	)
	good = np.flatnonzero(~upright[start:])
	recovery = round((int(good[-1]) + 1) * 0.02, 2) if len(good) and success else None
	return {
	"return": round(float(reward), 6),
	"success": success,
	"turn_deg": round(theta, 1),
	"max_tilt": round(inv, 1),
	"entry_mps": round(float((pos[start] - pos[start - 25]) / 0.5), 3),
	"entry_distance_m": round(float(pos[start] - pos[0]), 3),
	"final_m": round(glide, 3),
	"recovery_s": recovery,
	"inside": inside,
	}


	def test_job(args):
	theta, seed = args
	tr, _ = run(Landing(theta), seed)
	return metrics(tr)


	def calc(pool, models, seeds):
	res = list(
	pool.imap(test_job, [(v, s) for v in models for s in seeds], chunksize=2)
	)
	return [res[i : i + len(seeds)] for i in range(0, len(res), len(seeds))]


	def save(path, obj):
	tmp = path.with_suffix(".tmp")
	tmp.write_text(json.dumps(obj, indent=2) + "\n")
	tmp.replace(path)


	def train(
	steps,
	pop,
	batch,
	warm=None,
	run_dir=OUT,
	random_seed=11062026,
	validation=None,
	search_scale=0.62,
	floor=0.075,
	workers=16,
	):
	rng = np.random.default_rng(random_seed)
	mean = np.zeros(SIZE)
	sig = np.full(SIZE, 0.33)
	sig[24:] = 0.22
	best = []
	hist = []
	if warm is not None:
	mean = np.asarray(warm)
	sig *= search_scale
	with mp.Pool(
	workers, initializer=configure, initargs=(SPEEDS[1], TIMES[2])
	) as pool:
	validation = list(range(31, 43)) if validation is None else list(validation)
	baseline = calc(pool, [mean], validation)[0]
	print(
	"Baseline validation:",
	sum(r["success"] for r in baseline),
	"/",
	len(baseline),
	flush=True,
	)
	for gen in range(steps):
	first = rng.normal(size=(pop // 2, SIZE))
	eps = np.concatenate((first, -first))
	candidates = np.clip(mean + sig * eps, -2.8, 2.8)
	candidates[0] = mean
	seeds = list(
	map(int, rng.choice(np.arange(80, 600), size=batch, replace=False))
	)
	results = calc(pool, candidates, seeds)
	returns = np.array([np.mean([a["return"] for a in rr]) for rr in results])
	elite = candidates[np.argsort(returns)[-max(5, pop // 6) :]]
	mean = 0.27 * mean + 0.73 * np.mean(elite, axis=0)
	sig = np.maximum(floor, 0.62 * sig + 0.38 * np.std(elite, axis=0))
	leaders = [mean, candidates[np.argmax(returns)]]
	if (gen + 1) % 2 == 0 or gen == 0:
	trials = calc(pool, leaders, [*validation, 1])
	for model, rr in zip(leaders, trials):
	item = {
	"gen": gen + 1,
	"theta": model.tolist(),
	"validation_success": sum(q["success"] for q in rr[:-1]),
	"official_pass": rr[-1]["success"],
	"validation_return": float(
	np.mean([q["return"] for q in rr[:-1]])
	),
	}
	best.append(item)
	best.sort(
	key=lambda o: (o["validation_success"], o["validation_return"]),
	reverse=True,
	)
	best = best[:5] + [o for o in best[5:] if o["official_pass"]][:3]
	save(run_dir / "checkpoints.json", best)
	hist.append(
	{
	"generation": gen + 1,
	"training_seeds": seeds,
	"mean_return": float(returns.mean()),
	"best_training_return": float(returns.max()),
	}
	)
	save(run_dir / "learning.json", hist)
	print(
	"Generation",
	gen + 1,
	"best validation",
	best[0]["validation_success"],
	"/",
	len(validation),
	"training return",
	round(returns.max(), 2),
	flush=True,
	)