224 lines
9.9 KiB
Python
224 lines
9.9 KiB
Python
"""Optional installed-mjlab checks; GPU smoke is opt-in, never a full training run."""
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
SERVICE_ROOT = Path(__file__).resolve().parents[1]
|
|
for root in (SERVICE_ROOT, SERVICE_ROOT / "rl"):
|
|
sys.path.insert(0, str(root))
|
|
|
|
HAS_MJLAB = importlib.util.find_spec("mjlab") is not None
|
|
|
|
|
|
@unittest.skipUnless(HAS_MJLAB, "mjlab is not installed in this Python")
|
|
class ObstacleContractTest(unittest.TestCase):
|
|
def test_training_config_file_is_revalidated(self):
|
|
import json
|
|
import tempfile
|
|
|
|
from scripts.train import _load_task_config
|
|
from task_config import OBSTACLE_TASK, TaskConfigError, validate_task_config
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
source = Path(directory) / "training_config.json"
|
|
custom = validate_task_config(OBSTACLE_TASK, {"sensorCfg": {"fov": 60}}, 42)
|
|
source.write_text(json.dumps(custom))
|
|
self.assertEqual(_load_task_config(OBSTACLE_TASK, str(source), 42), custom)
|
|
with self.assertRaises(ValueError):
|
|
_load_task_config(OBSTACLE_TASK, str(source), 43)
|
|
custom["sensorCfg"]["maxDistance"] = -1
|
|
source.write_text(json.dumps(custom))
|
|
with self.assertRaises(TaskConfigError):
|
|
_load_task_config(OBSTACLE_TASK, str(source), 42)
|
|
source.write_text("x" * (128 * 1024 + 1))
|
|
with self.assertRaises(ValueError):
|
|
_load_task_config(OBSTACLE_TASK, str(source), 42)
|
|
|
|
def test_pattern_order_normalization_and_body_offsets(self):
|
|
import torch
|
|
from src.tasks.obstacle_avoidance.mdp import ForwardFanPatternCfg, forward_depth
|
|
|
|
offsets, directions = ForwardFanPatternCfg(fov=90).generate_rays(None, "cpu")
|
|
self.assertEqual(tuple(directions.shape), (32, 3))
|
|
torch.testing.assert_close(offsets, torch.tensor([0.3, 0, 0.05]).repeat(32, 1))
|
|
torch.testing.assert_close(directions.norm(dim=1), torch.ones(32))
|
|
self.assertAlmostEqual(directions[0, 1].item(), -(0.5**0.5), places=6)
|
|
self.assertAlmostEqual(directions[-1, 1].item(), 0.5**0.5, places=6)
|
|
scene = {
|
|
"forward_scan": SimpleNamespace(
|
|
data=SimpleNamespace(
|
|
distances=torch.tensor([[-1, 0, 1, 4, 8]], dtype=torch.float32)
|
|
)
|
|
)
|
|
}
|
|
result = forward_depth(SimpleNamespace(scene=scene), max_distance=4)
|
|
torch.testing.assert_close(result, torch.tensor([[1, 0, 0.25, 1, 1]]))
|
|
|
|
def test_navigation_matches_body_heading_and_arrival_stop(self):
|
|
import torch
|
|
from src.tasks.obstacle_avoidance.mdp import NavigationCommandCfg
|
|
|
|
data = SimpleNamespace(
|
|
root_link_pos_w=torch.tensor([[-5.0, 0, 0.32]]),
|
|
root_link_quat_w=torch.tensor([[1.0, 0, 0, 0]]),
|
|
)
|
|
|
|
class Scene(dict):
|
|
env_origins = torch.tensor([[-5.0, 0, 0]])
|
|
|
|
env = SimpleNamespace(
|
|
num_envs=1, device="cpu", scene=Scene(robot=SimpleNamespace(data=data))
|
|
)
|
|
command = NavigationCommandCfg(resampling_time_range=(1e9, 1e9)).build(env)
|
|
torch.testing.assert_close(command.command, torch.tensor([[0.6, 0, 0]]))
|
|
self.assertAlmostEqual(command.errors()[1].item(), 10)
|
|
data.root_link_quat_w[:] = torch.tensor([[0.5**0.5, 0, 0, 0.5**0.5]])
|
|
self.assertAlmostEqual(command.command[0, 2].item(), -1)
|
|
data.root_link_pos_w[0, 0] = 4.8
|
|
torch.testing.assert_close(command.command, torch.zeros((1, 3)))
|
|
|
|
def test_navigation_reset_randomizes_safe_distant_pairs_and_heading(self):
|
|
import torch
|
|
from src.tasks.obstacle_avoidance.mdp import NavigationCommandCfg
|
|
|
|
class Robot:
|
|
def __init__(self, count):
|
|
self.data = SimpleNamespace(
|
|
default_root_state=torch.tensor([[0, 0, 0.32] + [0] * 10] * count),
|
|
root_link_pos_w=torch.zeros((count, 3)),
|
|
root_link_quat_w=torch.zeros((count, 4)),
|
|
)
|
|
|
|
def write_root_link_pose_to_sim(self, pose, env_ids):
|
|
self.data.root_link_pos_w[env_ids] = pose[:, :3]
|
|
self.data.root_link_quat_w[env_ids] = pose[:, 3:]
|
|
|
|
def write_root_link_velocity_to_sim(self, velocity, env_ids):
|
|
self.velocity = velocity
|
|
|
|
class Scene(dict):
|
|
env_origins = torch.zeros((64, 3))
|
|
|
|
robot = Robot(64)
|
|
env = SimpleNamespace(num_envs=64, device="cpu", scene=Scene(robot=robot))
|
|
command = NavigationCommandCfg(
|
|
resampling_time_range=(1e9, 1e9),
|
|
navigation_points=((-2, 0), (-1, 0), (0, 0), (1, 0), (2, 0)),
|
|
component_starts=(0,),
|
|
component_counts=(5,),
|
|
fallback_pairs=(((-2, 0), (2, 0)),),
|
|
min_goal_distance=2,
|
|
).build(env)
|
|
torch.manual_seed(7)
|
|
command.sample_episode(torch.arange(64))
|
|
self.assertTrue(
|
|
((command.goals_w - robot.data.root_link_pos_w[:, :2]).norm(dim=1) >= 2).all()
|
|
)
|
|
self.assertGreater(torch.unique(robot.data.root_link_pos_w[:, :2], dim=0).shape[0], 1)
|
|
self.assertGreater(torch.unique(command.goals_w, dim=0).shape[0], 1)
|
|
torch.testing.assert_close(robot.data.root_link_quat_w.norm(dim=1), torch.ones(64))
|
|
self.assertTrue((robot.velocity == 0).all())
|
|
|
|
def test_all_terrain_presets_compile_at_exported_world_coordinates(self):
|
|
import mujoco
|
|
import numpy as np
|
|
from mjlab.terrains import TerrainGenerator
|
|
from src.tasks.obstacle_avoidance.env_cfg import unitree_go2_obstacle_env_cfg
|
|
from src.tasks.obstacle_avoidance.terrain import apply_terrain_configuration
|
|
from task_config import (
|
|
OBSTACLE_TASK,
|
|
TERRAIN_PRESETS,
|
|
build_terrain_layout,
|
|
validate_task_config,
|
|
)
|
|
|
|
for preset in (p for p in TERRAIN_PRESETS if p != "custom_boxes"):
|
|
with self.subTest(preset=preset):
|
|
custom = validate_task_config(OBSTACLE_TASK, {"terrainPreset": preset}, 7)
|
|
cfg = unitree_go2_obstacle_env_cfg()
|
|
apply_terrain_configuration(cfg, custom)
|
|
generator = TerrainGenerator(cfg.scene.terrain.terrain_generator)
|
|
spec = mujoco.MjSpec()
|
|
generator.compile(spec)
|
|
model = spec.compile()
|
|
layout = build_terrain_layout(custom)
|
|
np.testing.assert_allclose(model.geom_pos, [box["pos"] for box in layout["boxes"]])
|
|
np.testing.assert_allclose(
|
|
model.geom_size, [box["size"] for box in layout["boxes"]]
|
|
)
|
|
np.testing.assert_allclose(generator.terrain_origins[0, 0], [-5, 0, 0])
|
|
np.testing.assert_allclose(model.geom_friction[:, 0], layout["friction"])
|
|
self.assertTrue((model.geom_group == 0).all())
|
|
|
|
@unittest.skipUnless(os.environ.get("GO2_RUN_MJLAB_SMOKE") == "1", "opt-in GPU smoke")
|
|
def test_real_environment_81_observation_and_mujoco_raycast_parity(self):
|
|
import mujoco
|
|
import numpy as np
|
|
import torch
|
|
import warp as wp
|
|
from mjlab.envs import ManagerBasedRlEnv
|
|
from src.tasks.obstacle_avoidance.env_cfg import (
|
|
apply_obstacle_configuration,
|
|
unitree_go2_obstacle_env_cfg,
|
|
)
|
|
from src.tasks.obstacle_avoidance.mdp import ForwardFanPatternCfg
|
|
from task_config import JOINT_NAMES, OBSTACLE_TASK, validate_task_config
|
|
from warp._src import context
|
|
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("CUDA is unavailable")
|
|
if not hasattr(wp, "context"):
|
|
wp.context = context
|
|
cfg = unitree_go2_obstacle_env_cfg()
|
|
apply_obstacle_configuration(
|
|
cfg, validate_task_config(OBSTACLE_TASK, {}, 42), randomize_navigation=False
|
|
)
|
|
cfg.scene.num_envs = 2
|
|
env = ManagerBasedRlEnv(cfg, device="cuda:0")
|
|
try:
|
|
obs, _ = env.reset()
|
|
self.assertEqual(tuple(obs["actor"].shape), (2, 81))
|
|
self.assertEqual(list(env.scene["robot"].joint_names), JOINT_NAMES)
|
|
np.testing.assert_allclose(env.scene.env_origins.cpu(), [[-5, 0, 0]] * 2)
|
|
np.testing.assert_allclose(
|
|
env.scene["robot"].data.root_link_pos_w.cpu(), [[-5, 0, 0.32]] * 2, atol=1e-6
|
|
)
|
|
for _ in range(3):
|
|
obs, reward, terminated, truncated, _ = env.step(
|
|
torch.zeros((2, 12), device=env.device)
|
|
)
|
|
self.assertTrue(torch.isfinite(obs["actor"]).all())
|
|
self.assertTrue(torch.isfinite(reward).all())
|
|
self.assertFalse(terminated.any() or truncated.any())
|
|
model = env.sim.mj_model
|
|
data = mujoco.MjData(model)
|
|
data.qpos[:] = env.sim.data.qpos[0].cpu().numpy()
|
|
mujoco.mj_forward(model, data)
|
|
body = model.body("robot/base_link").id
|
|
rotation = data.xmat[body].reshape(3, 3)
|
|
offsets, directions = ForwardFanPatternCfg().generate_rays(None, "cpu")
|
|
expected = []
|
|
for offset, direction in zip(offsets.numpy(), directions.numpy(), strict=True):
|
|
distance = mujoco.mj_ray(
|
|
model,
|
|
data,
|
|
data.xpos[body] + rotation @ offset,
|
|
rotation @ direction,
|
|
np.array([1, 0, 0, 0, 0, 0], dtype=np.uint8),
|
|
1,
|
|
-1,
|
|
np.array([-1], dtype=np.int32),
|
|
)
|
|
expected.append(1 if distance < 0 else min(1, distance / 4))
|
|
np.testing.assert_allclose(obs["actor"][0, 47:79].cpu(), expected, atol=2e-5)
|
|
finally:
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|