513 lines
23 KiB
Python
513 lines
23 KiB
Python
"""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()
|