feat(training): release V0.9.1 避障训练与基础策略迁移
This commit is contained in:
@@ -0,0 +1,512 @@
|
||||
"""Obstacle-specific schema, orchestration, objective math and opt-in real rollout."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
SERVICE = Path(__file__).resolve().parents[1]
|
||||
for root in (SERVICE, SERVICE / "rl"):
|
||||
sys.path.insert(0, str(root))
|
||||
|
||||
from task_config import build_terrain_layout, validate_task_config # noqa: E402
|
||||
from tuning import obstacle_scoring as scoring # noqa: E402
|
||||
from tuning.advisor import DeepSeekAdvisor # noqa: E402
|
||||
from tuning.manager import TuningError, TuningManager # noqa: E402
|
||||
from tuning.process import GpuLease, ResourceBusyError # noqa: E402
|
||||
from tuning.schema import ( # noqa: E402
|
||||
OBSTACLE_TASK,
|
||||
RewardConfigError,
|
||||
apply_reward_configuration,
|
||||
base_configuration,
|
||||
merge_proposal,
|
||||
validate_configuration,
|
||||
validate_constraints,
|
||||
validate_proposal,
|
||||
)
|
||||
from tuning.scoring import EvaluationError # noqa: E402
|
||||
|
||||
|
||||
def sample(**kw):
|
||||
return (
|
||||
dict(
|
||||
distance=1.0,
|
||||
clearance=0.5,
|
||||
action_delta=0.0,
|
||||
ray_hit=0.0,
|
||||
collision=0.0,
|
||||
fall=0.0,
|
||||
terminal=0.0,
|
||||
**{},
|
||||
)
|
||||
| kw
|
||||
)
|
||||
|
||||
|
||||
def evaluation(custom, n=2):
|
||||
metrics = scoring.score_trajectory([sample()] * scoring.STEPS)
|
||||
return {
|
||||
"protocol": scoring.protocol(custom, n),
|
||||
"metrics": metrics,
|
||||
"seedMetrics": [
|
||||
{"seed": seed, "episodes": n, "rolloutSteps": scoring.STEPS, "metrics": metrics}
|
||||
for seed in scoring.SEEDS
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class ObstacleSchemaTest(unittest.TestCase):
|
||||
def test_ranges_task_isolation_nonfinite_and_patch_no_mutation(self):
|
||||
base = base_configuration(OBSTACLE_TASK)
|
||||
original = deepcopy(base)
|
||||
for section, key, low, high in (
|
||||
("weights", "avoidance_weight", 0.5, 5),
|
||||
("weights", "collision_penalty", -10, -0.5),
|
||||
("weights", "action_smoothness", -0.05, -0.001),
|
||||
("params", "target_velocity", 0.3, 1.2),
|
||||
):
|
||||
for value in (low, high):
|
||||
changed = deepcopy(base)
|
||||
changed[section][key] = value
|
||||
self.assertEqual(validate_configuration(changed, OBSTACLE_TASK), changed)
|
||||
for value in (low - 0.00001, high + 0.00001, float("nan"), float("inf"), True):
|
||||
with self.subTest(key=key, value=value), self.assertRaises(RewardConfigError):
|
||||
validate_proposal({section: {key: value}}, base, task_id=OBSTACLE_TASK)
|
||||
for proposal in (
|
||||
{"weights": {"pose": 1.1}},
|
||||
{"weights": {"target_velocity": 0.8}},
|
||||
{"params": {"seed": 101}},
|
||||
{"unknown": {}},
|
||||
{"weights": {"avoidance_weight": 4.01}},
|
||||
):
|
||||
with self.assertRaises(RewardConfigError):
|
||||
validate_proposal(proposal, base, task_id=OBSTACLE_TASK)
|
||||
with self.assertRaises(RewardConfigError):
|
||||
validate_configuration(base)
|
||||
with self.assertRaises(RewardConfigError):
|
||||
validate_configuration(base_configuration(), OBSTACLE_TASK)
|
||||
with self.assertRaises(RewardConfigError):
|
||||
validate_constraints({"weights.pose": {"kind": "fixed", "value": 1}}, OBSTACLE_TASK)
|
||||
candidate = merge_proposal(
|
||||
base, {"params": {"target_velocity": 1.2}}, task_id=OBSTACLE_TASK
|
||||
)
|
||||
self.assertEqual(candidate["params"]["target_velocity"], 1.2)
|
||||
self.assertEqual(base, original)
|
||||
|
||||
def test_real_training_config_mapping_and_command(self):
|
||||
import torch
|
||||
from scripts.train import (
|
||||
_configure_task_and_rewards,
|
||||
_load_reward_config,
|
||||
_load_task_config,
|
||||
)
|
||||
from src.tasks.obstacle_avoidance.env_cfg import (
|
||||
apply_obstacle_configuration,
|
||||
unitree_go2_obstacle_env_cfg,
|
||||
)
|
||||
|
||||
base = base_configuration(OBSTACLE_TASK)
|
||||
base["weights"].update(avoidance_weight=3, collision_penalty=-7, action_smoothness=-0.02)
|
||||
base["params"]["target_velocity"] = 0.9
|
||||
custom = validate_task_config(OBSTACLE_TASK, {"sensorCfg": {"fov": 60}}, 42)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
reward, task = Path(directory) / "reward.json", Path(directory) / "task.json"
|
||||
reward.write_text(json.dumps(base))
|
||||
task.write_text(json.dumps(custom))
|
||||
cfg = unitree_go2_obstacle_env_cfg()
|
||||
apply_obstacle_configuration(cfg, _load_task_config(OBSTACLE_TASK, str(task), 42))
|
||||
apply_reward_configuration(
|
||||
cfg, _load_reward_config(str(reward), None, OBSTACLE_TASK), OBSTACLE_TASK
|
||||
)
|
||||
deployment, checked = _configure_task_and_rewards(
|
||||
OBSTACLE_TASK,
|
||||
SimpleNamespace(
|
||||
env=unitree_go2_obstacle_env_cfg(),
|
||||
task_config=str(task),
|
||||
reward_config=str(reward),
|
||||
reward_config_json=None,
|
||||
agent=SimpleNamespace(seed=42),
|
||||
),
|
||||
)
|
||||
self.assertEqual(checked, base)
|
||||
self.assertEqual(deployment["navigation"]["speed"], 0.9)
|
||||
self.assertEqual(deployment["sensorCfg"]["avoidanceWeight"], 3)
|
||||
self.assertEqual(cfg.rewards["obstacle_proximity"].weight, -3)
|
||||
self.assertEqual(cfg.rewards["obstacle_collision"].weight, -7)
|
||||
self.assertIs(
|
||||
cfg.rewards["obstacle_collision"].func, cfg.terminations["illegal_contact"].func
|
||||
)
|
||||
self.assertEqual(cfg.rewards["action_rate_l2"].weight, -0.02)
|
||||
self.assertNotIn("target_velocity", cfg.rewards)
|
||||
self.assertEqual(cfg.commands["twist"].speed, 0.9)
|
||||
|
||||
class Scene(dict):
|
||||
env_origins = torch.tensor([[-5.0, 0, 0]])
|
||||
|
||||
scene = Scene(
|
||||
robot=SimpleNamespace(
|
||||
data=SimpleNamespace(
|
||||
root_link_pos_w=torch.tensor([[-5.0, 0, 0.32]]),
|
||||
root_link_quat_w=torch.tensor([[1.0, 0, 0, 0]]),
|
||||
)
|
||||
)
|
||||
)
|
||||
command = cfg.commands["twist"].build(
|
||||
SimpleNamespace(num_envs=1, device="cpu", scene=scene)
|
||||
)
|
||||
self.assertAlmostEqual(command.command[0, 0].item(), 0.9, places=6)
|
||||
|
||||
def test_mock_deepseek_task_context_no_network(self):
|
||||
advisor = DeepSeekAdvisor()
|
||||
output = SimpleNamespace(
|
||||
weights={"avoidance_weight": 2.2},
|
||||
params={"target_velocity": 0.7},
|
||||
rationale="绕行与目标导航",
|
||||
expected_impact={},
|
||||
confidence=0.8,
|
||||
)
|
||||
fake = SimpleNamespace(run_sync=lambda prompt: SimpleNamespace(output=output))
|
||||
with patch.object(advisor, "_agent", return_value=fake):
|
||||
result = advisor.propose({"task": OBSTACLE_TASK}, base_configuration(OBSTACLE_TASK))
|
||||
self.assertEqual(result["patch"]["params"], {"target_velocity": 0.7})
|
||||
|
||||
|
||||
class ObstacleScoringTest(unittest.TestCase):
|
||||
def test_hand_calculated_scores_and_first_terminal(self):
|
||||
samples = [
|
||||
sample(clearance=0.25, action_delta=0.5),
|
||||
sample(distance=0.4, clearance=0.25, action_delta=0.5),
|
||||
]
|
||||
metrics = scoring.score_trajectory(samples, 2)
|
||||
self.assertEqual(metrics["success"], 1)
|
||||
self.assertEqual(metrics["time"], 0)
|
||||
self.assertEqual(metrics["clearance"], 0.5)
|
||||
self.assertEqual(metrics["smooth"], 0.5)
|
||||
self.assertAlmostEqual(scoring.score_evaluation(metrics)["score"], 0.65)
|
||||
failed = scoring.score_trajectory([sample(distance=0.1, terminal=1, fall=1)], 1000)
|
||||
self.assertEqual((failed["success"], failed["time"], failed["clearance"]), (0, 0, 0))
|
||||
contact = scoring.score_trajectory([sample(distance=0.1, terminal=1, collision=1)], 1000)
|
||||
self.assertEqual((contact["arrival_rate"], contact["success"], contact["time"]), (1, 0, 0))
|
||||
with self.assertRaises(EvaluationError):
|
||||
scoring.score_trajectory([sample()], 1000)
|
||||
with self.assertRaises(EvaluationError):
|
||||
scoring.score_trajectory([sample(terminal=1), sample()], 2)
|
||||
with self.assertRaises(EvaluationError):
|
||||
scoring.score_trajectory([sample(clearance=float("nan"))], 1)
|
||||
|
||||
def test_recorder_excludes_floor_and_keeps_terminal_before_reset(self):
|
||||
import torch
|
||||
from scripts.evaluate_obstacle import FirstEpisodeRecorder
|
||||
|
||||
class Scene(dict):
|
||||
env_origins = torch.zeros((2, 3))
|
||||
|
||||
robot = SimpleNamespace(
|
||||
root_link_pos_w=torch.tensor([[0.0, 0, 0.32], [0.0, 0, 0.1]]),
|
||||
projected_gravity_b=torch.tensor([[0.0, 0, -1.0], [0.0, 1.0, 0.0]]),
|
||||
)
|
||||
scene = Scene(
|
||||
robot=SimpleNamespace(data=robot),
|
||||
nonfoot_ground_touch=SimpleNamespace(
|
||||
data=SimpleNamespace(force_history=torch.zeros((2, 1, 4, 3)))
|
||||
),
|
||||
forward_scan=SimpleNamespace(
|
||||
data=SimpleNamespace(distances=torch.ones((2, 32))),
|
||||
cfg=SimpleNamespace(max_distance=4),
|
||||
),
|
||||
)
|
||||
command = SimpleNamespace(errors=lambda: (None, torch.tensor([0.2, 0.2]), None))
|
||||
env = SimpleNamespace(
|
||||
scene=scene,
|
||||
device="cpu",
|
||||
num_envs=2,
|
||||
action_manager=SimpleNamespace(action=torch.ones((2, 12)) * 0.5),
|
||||
command_manager=SimpleNamespace(get_term=lambda name: command),
|
||||
termination_manager=SimpleNamespace(compute=lambda: torch.ones(2, dtype=torch.bool)),
|
||||
)
|
||||
floor = {"pos": [0, 0, -0.1], "size": [6, 6, 0.1]}
|
||||
recorder = FirstEpisodeRecorder(env, {"spawn": [0, 0, 0.32], "boxes": [floor]})
|
||||
env.termination_manager.compute()
|
||||
robot.root_link_pos_w[:] = torch.tensor([0.0, 0, 0.32]) # Simulated auto-reset overwrite.
|
||||
robot.projected_gravity_b[:] = torch.tensor([0.0, 0, -1.0])
|
||||
env.termination_manager.compute()
|
||||
self.assertEqual([len(s) for s in recorder.samples], [1, 1])
|
||||
upright, lying = [scoring.score_trajectory(s, 2) for s in recorder.samples]
|
||||
self.assertEqual(upright["clearance"], 0.5) # 1 valid safe step / 2, not floor distance 0.
|
||||
self.assertEqual(upright["success"], 1)
|
||||
self.assertEqual(lying["clearance"], 0)
|
||||
self.assertEqual(lying["success"], 0)
|
||||
self.assertEqual(lying["time"], 0)
|
||||
self.assertEqual(lying["fall_rate"], 1)
|
||||
|
||||
def test_threshold_boundary_and_fail_closed(self):
|
||||
base = scoring.score_trajectory([sample(distance=0.1)], 1)
|
||||
base.update(success=0.5, fall_rate=0.1, no_fall=0.9)
|
||||
current = dict(base, success=0.48, fall_rate=0.12, no_fall=0.88)
|
||||
self.assertTrue(scoring.score_evaluation(current, base)["eligible"])
|
||||
for changed in (
|
||||
dict(current, success=0.479999),
|
||||
dict(current, fall_rate=0.120001, no_fall=0.879999),
|
||||
):
|
||||
self.assertFalse(scoring.score_evaluation(changed, base)["eligible"])
|
||||
self.assertEqual(scoring.score_evaluation(changed, base)["score"], -1)
|
||||
custom = validate_task_config(OBSTACLE_TASK, {}, 42)
|
||||
expected = scoring.protocol(custom, 2)
|
||||
valid = evaluation(custom)
|
||||
scoring.validate_evaluation(valid, expected)
|
||||
for mutate in (
|
||||
lambda v: v["seedMetrics"].pop(),
|
||||
lambda v: v["metrics"].pop("success"),
|
||||
lambda v: v["protocol"].update(stepsPerSeed=999),
|
||||
lambda v: v["seedMetrics"][0].update(episodes=1),
|
||||
lambda v: v["metrics"].update(smooth=float("inf")),
|
||||
):
|
||||
bad = deepcopy(valid)
|
||||
mutate(bad)
|
||||
with self.assertRaises(EvaluationError):
|
||||
scoring.validate_evaluation(bad, expected)
|
||||
|
||||
def test_fixed_three_seed_and_custom_authoritative_map(self):
|
||||
custom = validate_task_config(OBSTACLE_TASK, {}, 42)
|
||||
scenes = scoring.evaluation_scenarios(custom)
|
||||
self.assertEqual([s["seed"] for s in scenes], [101, 202, 303])
|
||||
self.assertNotEqual(scenes[0]["terrain"], scenes[1]["terrain"])
|
||||
layout = build_terrain_layout(custom)
|
||||
layout["approximation"] = True
|
||||
custom = validate_task_config(
|
||||
OBSTACLE_TASK, {"terrainPreset": "custom_boxes", "customTerrainBoxes": layout}, 42
|
||||
)
|
||||
fixed = scoring.protocol(custom, 2)
|
||||
self.assertEqual(fixed["sceneMode"], "fixed-custom-map")
|
||||
self.assertEqual([s["terrain"] for s in fixed["scenarios"]], [layout] * 3)
|
||||
self.assertEqual(fixed, scoring.protocol(deepcopy(custom), 2))
|
||||
|
||||
|
||||
class ObstacleManagerTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
root = Path(self.temp.name)
|
||||
self.manager = TuningManager(SERVICE / "rl", sys.executable, root, GpuLease())
|
||||
payload = {
|
||||
"taskId": OBSTACLE_TASK,
|
||||
"mode": "approval",
|
||||
"evalNumEnvs": 2,
|
||||
"numEnvs": 2,
|
||||
"initialIterations": 1,
|
||||
"middleIterations": 1,
|
||||
"finalIterations": 1,
|
||||
"sensorCfg": {"fov": 60, "avoidanceWeight": 3},
|
||||
}
|
||||
mode, config, objective, fallback = self.manager.parse_create(payload)
|
||||
self.session = self.manager.storage.create_session(mode, config, objective, fallback)
|
||||
self.trial = self.manager.storage.create_trial(
|
||||
self.session["id"],
|
||||
0,
|
||||
0,
|
||||
1,
|
||||
self.manager._base_configuration(self.session),
|
||||
None,
|
||||
"trial-000-rung-0",
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
self.manager.shutdown()
|
||||
self.temp.cleanup()
|
||||
|
||||
def test_argv_json_real_train_and_eval_ingress(self):
|
||||
commands = []
|
||||
custom = self.session["config"]["taskConfig"]
|
||||
|
||||
def run(session_id, command, cwd, environment, log_path):
|
||||
commands.append(command)
|
||||
task_path = Path(command[command.index("--task-config") + 1])
|
||||
reward_path = Path(command[command.index("--reward-config") + 1])
|
||||
from scripts.train import _load_reward_config, _load_task_config
|
||||
|
||||
self.assertEqual(_load_task_config(OBSTACLE_TASK, str(task_path), 42), custom)
|
||||
self.assertEqual(
|
||||
_load_reward_config(str(reward_path), None, OBSTACLE_TASK),
|
||||
self.trial["rewardConfig"],
|
||||
)
|
||||
if "scripts/train.py" in command:
|
||||
(log_path.parent / "model_0.pt").write_bytes(b"mock-checkpoint")
|
||||
(log_path.parent / "policy.onnx").write_bytes(b"mock-onnx")
|
||||
else:
|
||||
Path(command[command.index("--output") + 1]).write_text(
|
||||
json.dumps(evaluation(custom))
|
||||
)
|
||||
return 0
|
||||
|
||||
self.manager.cancel_events[self.session["id"]] = threading.Event()
|
||||
with (
|
||||
patch.object(self.manager, "_run_command", side_effect=run),
|
||||
patch.object(self.manager.studies, "record", return_value=0),
|
||||
):
|
||||
result = self.manager._execute_trial(self.session, self.trial)
|
||||
self.assertEqual(result["state"], "completed")
|
||||
self.assertIn("--steps-per-seed=1000", commands[1])
|
||||
self.assertIn(OBSTACLE_TASK, commands[0])
|
||||
context = self.manager._proposal_context(self.session)
|
||||
self.assertIn("32", context["taskContext"])
|
||||
self.assertEqual(set(context["allowlist"]["params"]), {"target_velocity"})
|
||||
|
||||
def test_cas_stale_guardrails_and_mode_roundtrip(self):
|
||||
sid = self.session["id"]
|
||||
constraint = {"params.target_velocity": {"kind": "range", "min": 0.4, "max": 0.8}}
|
||||
self.manager.set_constraints(sid, {"revision": 0, "constraints": constraint})
|
||||
before = self.manager.detail(sid)
|
||||
with self.assertRaises(ResourceBusyError):
|
||||
self.manager.set_constraints(sid, {"revision": 0, "constraints": {}})
|
||||
self.assertEqual(self.manager.detail(sid), before)
|
||||
for bad in (
|
||||
{"weights.pose": {"kind": "fixed", "value": 1}},
|
||||
{"params.target_velocity": {"kind": "fixed", "value": float("nan")}},
|
||||
):
|
||||
with self.assertRaises(RewardConfigError):
|
||||
self.manager.set_constraints(sid, {"revision": 1, "constraints": bad})
|
||||
self.assertEqual(self.manager.detail(sid), before)
|
||||
self.manager.storage.update_session(sid, state="awaiting_approval")
|
||||
proposal = self.manager.storage.create_proposal(
|
||||
sid, self.trial["id"], {"params": {"target_velocity": 0.7}}, "test", {}, 0.8, "agent"
|
||||
)
|
||||
self.assertEqual(self.manager.set_mode(sid, {"mode": "automatic"})["mode"], "automatic")
|
||||
self.assertEqual(self.manager.storage.get_proposal(proposal["id"])["state"], "approved")
|
||||
self.assertEqual(self.manager.set_mode(sid, {"mode": "approval"})["mode"], "approval")
|
||||
with self.assertRaises(RewardConfigError):
|
||||
validate_proposal(
|
||||
{"params": {"target_velocity": 0.9}},
|
||||
self.trial["rewardConfig"],
|
||||
constraint,
|
||||
OBSTACLE_TASK,
|
||||
)
|
||||
|
||||
def test_protocol_fields_cannot_be_changed(self):
|
||||
for values in (
|
||||
{"evalSteps": 999},
|
||||
{"seeds": [1, 2, 3]},
|
||||
{"objectiveWeights": {}},
|
||||
{"taskConfig": {"unknown": 1}},
|
||||
{"taskConfig": {"seed": True}},
|
||||
{"seed": 1, "taskConfig": {"seed": True}},
|
||||
):
|
||||
with self.assertRaises(TuningError):
|
||||
self.manager.parse_create({"taskId": OBSTACLE_TASK, **values})
|
||||
|
||||
|
||||
class IsolatedEvaluationTest(unittest.TestCase):
|
||||
def test_spawn_results_and_failure_cancel_missing_seed_fail_closed(self):
|
||||
import subprocess
|
||||
|
||||
from scripts.evaluate import EvaluateConfig
|
||||
from scripts.evaluate_obstacle import evaluate_isolated_seeds
|
||||
|
||||
custom = validate_task_config(OBSTACLE_TASK, {}, 42)
|
||||
fixed = scoring.protocol(custom, 2)
|
||||
reward = base_configuration(OBSTACLE_TASK)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
checkpoint = Path(directory) / "model.pt"
|
||||
checkpoint.write_bytes(b"test checkpoint snapshot")
|
||||
cfg = EvaluateConfig(checkpoint=str(checkpoint), output="unused", num_envs=2)
|
||||
|
||||
def run(command, **kwargs):
|
||||
self.assertFalse(kwargs["shell"])
|
||||
self.assertEqual(command[0], sys.executable)
|
||||
request = json.loads(Path(command[2]).read_text())
|
||||
result = evaluation(custom)["seedMetrics"][0]
|
||||
result["seed"] = request["scenario"]["seed"]
|
||||
result["checkpointSha256"] = request["checkpointSha256"]
|
||||
Path(command[3]).write_text(json.dumps(result))
|
||||
|
||||
with patch("scripts.evaluate_obstacle.subprocess.run", side_effect=run) as mocked:
|
||||
results = evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward)
|
||||
self.assertEqual([r["seed"] for r in results], [101, 202, 303])
|
||||
self.assertEqual(mocked.call_count, 3)
|
||||
for failure in (
|
||||
subprocess.CalledProcessError(1, "worker"),
|
||||
subprocess.TimeoutExpired("worker", 1200),
|
||||
KeyboardInterrupt(),
|
||||
):
|
||||
with patch(
|
||||
"scripts.evaluate_obstacle.subprocess.run", side_effect=failure
|
||||
) as mocked:
|
||||
with self.assertRaises(type(failure)):
|
||||
evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward)
|
||||
self.assertEqual(mocked.call_count, 1)
|
||||
with (
|
||||
patch("scripts.evaluate_obstacle.subprocess.run"),
|
||||
self.assertRaises(FileNotFoundError),
|
||||
):
|
||||
evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward)
|
||||
|
||||
def wrong_seed(command, **kwargs):
|
||||
run(command, **kwargs)
|
||||
path = Path(command[3])
|
||||
result = json.loads(path.read_text())
|
||||
result["seed"] = 0
|
||||
path.write_text(json.dumps(result))
|
||||
|
||||
with (
|
||||
patch("scripts.evaluate_obstacle.subprocess.run", side_effect=wrong_seed),
|
||||
self.assertRaises(EvaluationError),
|
||||
):
|
||||
evaluate_isolated_seeds(OBSTACLE_TASK, cfg, fixed, reward)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
os.environ.get("GO2_RUN_TUNING_SMOKE") == "1", "opt-in real GPU checkpoint rollout"
|
||||
)
|
||||
class ObstacleRolloutSmoke(unittest.TestCase):
|
||||
def test_actual_checkpoint_policy_statistics_and_first_terminal(self):
|
||||
from dataclasses import asdict
|
||||
|
||||
import src.tasks # noqa: F401
|
||||
import torch
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from mjlab.rl import RslRlVecEnvWrapper
|
||||
from mjlab.tasks.registry import load_rl_cfg, load_runner_cls
|
||||
from scripts.evaluate import EvaluateConfig
|
||||
from scripts.evaluate_obstacle import FirstEpisodeRecorder, configure_seed
|
||||
|
||||
cfg = EvaluateConfig(
|
||||
checkpoint=os.environ["GO2_TUNING_CHECKPOINT"],
|
||||
output="/tmp/unused.json",
|
||||
num_envs=2,
|
||||
device="cuda:0",
|
||||
)
|
||||
custom = validate_task_config(OBSTACLE_TASK, {}, 42)
|
||||
scenario = scoring.evaluation_scenarios(custom)[0]
|
||||
env_cfg = configure_seed(OBSTACLE_TASK, cfg, scenario, base_configuration(OBSTACLE_TASK))
|
||||
env_cfg.episode_length_s = 0.02 # Test-only one-step timeout exercises auto-reset capture.
|
||||
env = ManagerBasedRlEnv(env_cfg, device="cuda:0")
|
||||
agent_cfg = load_rl_cfg(OBSTACLE_TASK)
|
||||
wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
try:
|
||||
runner = load_runner_cls(OBSTACLE_TASK)(
|
||||
wrapped, asdict(agent_cfg), log_dir=None, device="cuda:0"
|
||||
)
|
||||
runner.load(
|
||||
cfg.checkpoint, load_cfg={"actor": True}, strict=True, map_location="cuda:0"
|
||||
)
|
||||
saved = torch.load(cfg.checkpoint, weights_only=False)["actor_state_dict"]
|
||||
actual = runner.alg.actor.state_dict()
|
||||
normalizers = [k for k in saved if "normaliz" in k]
|
||||
self.assertTrue(normalizers)
|
||||
for key in normalizers:
|
||||
torch.testing.assert_close(actual[key], saved[key].to(actual[key].device))
|
||||
obs, _ = env.reset(seed=101)
|
||||
recorder = FirstEpisodeRecorder(env, scenario["terrain"])
|
||||
with torch.inference_mode():
|
||||
policy = runner.get_inference_policy(device="cuda:0")
|
||||
for _ in range(2):
|
||||
obs, _, _, _ = wrapped.step(policy(obs))
|
||||
self.assertEqual([len(s) for s in recorder.samples], [1, 1])
|
||||
self.assertTrue(all(s[0]["terminal"] for s in recorder.samples))
|
||||
scoring.validate_metrics(recorder.metrics(horizon=2))
|
||||
finally:
|
||||
wrapped.close()
|
||||
Reference in New Issue
Block a user