File size: 94,427 Bytes
a1d2967 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 1710 1711 1712 1713 1714 1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 | """
embodied_pcei/src/ninja_env.py
High-performance physical simulation environment for the American Ninja Warrior course.
Built on MuJoCo with native Apple Silicon performance.
Features:
- Articulated bipedal humanoid with 21 joint torque actuators.
- Staged American Ninja Warrior course:
* Stage 0: Flat sprint runway (0 to 8m)
* Stage 1: Quad steps / angled stepping stones (8 to 18m)
* Stage 2: Elevated balance beam (18 to 28m)
* Stage 3: Hurdles & stepping barriers (28 to 38m)
* Stage 4: Warped wall incline ramp & summit buzzer (38 to 48m)
- Raycast Pseudo-Vision: forward height probes measuring terrain/obstacle topology.
- Metabolic / Homeostatic Telemetry: lungs/stamina, stomach/hunger, heart/health.
- Built-in offscreen renderer for live web dashboard streaming.
"""
import numpy as np
import mujoco
from pathlib import Path
from typing import Dict, Tuple, Optional
XML_PATH = Path(__file__).resolve().parent.parent / "models" / "ninja_course.xml"
FLAT_ARENA_XML = Path(__file__).resolve().parent.parent / "models" / "flat_arena.xml"
# Obstacle Stage Milestones along +X
STAGE_MILESTONES = [
(8.0, "Stage 1: Quad Steps"),
(18.0, "Stage 2: Balance Beam"),
(28.0, "Stage 3: Hurdles"),
(38.0, "Stage 4: Warped Wall"),
(46.0, "Course Finish: Summit Buzzer"),
]
# 18 Progressive Navigation Waypoints across the 48m Ninja Warrior Course
# Format: (Target 3D Coordinate [X, Y, Z], Waypoint Label, Arrival Radius in meters)
COURSE_WAYPOINTS = [
(np.array([3.0, 0.00, 1.33], dtype=np.float32), "Runway Accelerate 1", 1.2),
(np.array([6.0, 0.00, 1.33], dtype=np.float32), "Runway Accelerate 2", 1.2),
(np.array([8.0, 0.00, 1.33], dtype=np.float32), "Runway Exit Threshold", 1.0),
(np.array([10.0, -0.60, 0.20], dtype=np.float32), "Stage 1: Quad Step 1 (Right)", 0.8),
(np.array([12.2, 0.60, 0.30], dtype=np.float32), "Stage 1: Quad Step 2 (Left)", 0.8),
(np.array([14.4, -0.50, 0.35], dtype=np.float32), "Stage 1: Quad Step 3 (Right)", 0.8),
(np.array([16.6, 0.50, 0.40], dtype=np.float32), "Stage 1: Quad Step 4 (Left)", 0.8),
(np.array([18.5, 0.00, 0.45], dtype=np.float32), "Platform Mid 1", 1.0),
(np.array([20.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Entry", 0.6),
(np.array([23.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Mid", 0.6),
(np.array([26.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Exit", 0.6),
(np.array([28.5, 0.00, 0.50], dtype=np.float32), "Platform Mid 2", 1.0),
(np.array([31.5, 0.00, 0.20], dtype=np.float32), "Stage 3: Hurdle 1", 0.9),
(np.array([34.5, 0.00, 0.28], dtype=np.float32), "Stage 3: Hurdle 2", 0.9),
(np.array([37.5, 0.00, 0.22], dtype=np.float32), "Stage 3: Hurdle 3", 0.9),
(np.array([40.0, 0.00, 0.30], dtype=np.float32), "Stage 4: Warped Wall Approach",1.2),
(np.array([42.5, 0.00, 0.70], dtype=np.float32), "Stage 4: Warped Wall Incline", 1.2),
(np.array([47.5, 0.00, 2.45], dtype=np.float32), "Course Finish: Summit Buzzer", 0.6),
]
# Translucent Neon Color Palette for Multi-Agent Ghost Swarm Rendering
GHOST_PALETTE = [
np.array([0.2, 0.7, 1.0, 0.16], dtype=np.float32), # Neon Cyan
np.array([0.7, 0.3, 1.0, 0.16], dtype=np.float32), # Electric Purple
np.array([0.1, 0.9, 0.6, 0.16], dtype=np.float32), # Neon Emerald
np.array([1.0, 0.55, 0.1, 0.16], dtype=np.float32), # Sunset Orange
np.array([1.0, 0.2, 0.6, 0.16], dtype=np.float32), # Neon Pink
np.array([0.35, 0.8, 1.0, 0.16], dtype=np.float32), # Sky Blue
np.array([0.85, 0.95, 0.2, 0.16], dtype=np.float32),# Lime Volt
]
# Canonical Reference Quiet Standing Equilibrium Pose (21 Actuated Joints)
# Torso at Z=1.168m, spine vertical, knees naturally flexed at 0.05 rad (~3 deg), arms down at sides
STAND_REF_QPOS = np.array([
0.0, 0.0, 0.0, # abdomen_z, abdomen_y, abdomen_x
0.0, 0.0, -0.05, 0.05, -0.025, 0.0, # right_hip_x, z, y, knee, ankle_y, ankle_x
0.0, 0.0, -0.05, 0.05, -0.025, 0.0, # left_hip_x, z, y, knee, ankle_y, ankle_x
0.65, 0.0, 0.0, # right_shoulder_x, y, elbow (down at side)
-0.65, 0.0, 0.0, # left_shoulder_x, y, elbow (down at side)
], dtype=np.float32)
STAND_JOINT_WEIGHTS = np.array([
1.5, 1.5, 1.5, # abdomen_z, abdomen_y, abdomen_x
1.5, 1.5, 1.5, 1.5, 1.5, 1.5, # right leg (hip_x, hip_z, hip_y, knee, ankle_y, ankle_x)
1.5, 1.5, 1.5, 1.5, 1.5, 1.5, # left leg (hip_x, hip_z, hip_y, knee, ankle_y, ankle_x)
0.5, 0.5, 0.5, # right arm (shoulder_x, shoulder_y, elbow)
0.5, 0.5, 0.5 # left arm (shoulder_x, shoulder_y, elbow)
], dtype=np.float32)
class NinjaEnv:
def __init__(self, xml_path: Optional[str] = None, frame_skip: int = 5, arena_mode: bool = True, stream_mode: str = "cruise"):
self.arena_mode = arena_mode
self.stream_mode = stream_mode # "cruise", "stand", "getup", "reach"
if xml_path is None:
xml_path = str(FLAT_ARENA_XML if arena_mode else XML_PATH)
self.model = mujoco.MjModel.from_xml_path(xml_path)
self.data = mujoco.MjData(self.model)
self.frame_skip = frame_skip
self.dt = self.model.opt.timestep * frame_skip # 0.005 * 5 = 0.025s (40 Hz)
self.n_actuators = self.model.nu # 21
self.n_qpos = self.model.nq # 28 (7 root + 21 joints)
self.n_qvel = self.model.nv # 27 (6 root + 21 joints)
# Solid-State Angular LiDAR Sensor Configuration
# 16 vertical beams (pitch: -35 deg to +15 deg) x 32 horizontal beams (yaw: -45 deg to +45 deg) = 512 rays
self.num_lidar_v = 16
self.num_lidar_h = 32
self.num_lidar_rays = self.num_lidar_v * self.num_lidar_h # 512
self.num_rays = self.num_lidar_rays # backwards-compatibility alias
self.lidar_max_range = 10.0 # meters
self.lidar_v_angles = np.radians(np.linspace(-35, 15, self.num_lidar_v))
self.lidar_h_angles = np.radians(np.linspace(-45, 45, self.num_lidar_h))
# Precompute unit direction vectors in head local coordinate frame
# Forward = +X, Left = +Y, Up = +Z
local_vectors = []
for pitch in self.lidar_v_angles:
for yaw in self.lidar_h_angles:
vx = np.cos(pitch) * np.cos(yaw)
vy = np.cos(pitch) * np.sin(yaw)
vz = np.sin(pitch)
v = np.array([vx, vy, vz], dtype=np.float64)
v /= np.linalg.norm(v)
local_vectors.append(v)
self.lidar_local_vectors = np.array(local_vectors, dtype=np.float64) # (128, 3)
self.current_lidar_grid = np.zeros((self.num_lidar_v, self.num_lidar_h), dtype=np.float32)
self._geomid_buf = np.zeros(1, dtype=np.int32)
self._visor_offset = np.array([0.15, 0.0, 0.0], dtype=np.float64)
# Metabolic internal states
self.stamina = 1.0 # 0.0 to 1.0
self.health = 1.0 # 0.0 to 1.0
self.hunger = 0.0 # 0.0 to 1.0
# Milestone tracking & Course Waypoints
self.milestones_reached = set()
self.waypoints_reached = set()
self.current_wp_idx = 0
self.furthest_x = 0.0
self.max_steps = 1000
# Flat Arena & Dynamic Random Goal Navigation State (Stream 4: Cruise)
self.goal_pos = np.array([5.0, 0.0, 1.25], dtype=np.float32)
self.goals_completed = 0
self.goal_arrival_radius = 2.0 # Wide capture radius (2.0m): easily clears by walking near target
self.min_dist_to_goal = 5.0
self.prev_dist_to_goal = 5.0
self.steps_not_approaching = 0
self.transition_leeway_window = 90 # ~2.25s leeway (90 steps at 40Hz) for turning & realignment
self.goal_transition_steps = 0 # No leeway on initial spawn (must advance immediately)
# Stream 1: Quiet Standing Stillness State
self.stand_still_steps = 0
self.stand_drift_steps = 0
self.stand_one_leg_steps = 0
# Stream 2: Fall Recovery State (Ragdoll & Curated Library)
self.fall_state_type = "random_ragdoll"
self.consecutive_upright_steps = 0
self.prev_z = 0.20
self.prev_head_z = 0.20
self.getup_stagnant_steps = 0
self.max_elevation_achieved = 0.20
self.max_head_elevation = 0.20
self.getup_goal_pos = np.array([0.0, 0.0, 1.168], dtype=np.float32)
self.getup_success_count = 0
# Stream 3: Bilateral Hand Reaching State (Bolted Torso)
self.active_reach_hand = "right"
self.reach_target_pos = np.array([0.35, -0.25, 1.35], dtype=np.float32)
self.reach_touches = 0
self.prev_reach_dist = 0.40
self.min_reach_dist = 0.40
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
# Bilateral gait stance duty-cycle counters
self.r_foot_contact_steps = 0
self.l_foot_contact_steps = 0
# Offscreen renderer instance (lazy-loaded)
self._renderer = None
# Cache body and geom IDs for fast lookups
self.torso_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "torso")
self.head_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "head")
self.r_foot_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "right_foot")
self.l_foot_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "left_foot")
self.r_foot_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_foot_geom")
self.l_foot_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_foot_geom")
self.r_hand_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_hand")
self.l_hand_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_hand")
self.r_arm_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "right_lower_arm")
self.l_arm_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "left_lower_arm")
# Cache leg joint qpos indices for athletic stance and gait monitoring
self.r_knee_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_knee")]
self.l_knee_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_knee")]
self.r_hip_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_hip_y")]
self.l_hip_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_hip_y")]
self.r_ankle_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_ankle_y")]
self.l_ankle_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_ankle_y")]
# Cache abdomen joint qpos indices for core posture alignment
self.abdomen_x_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_x")]
self.abdomen_y_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_y")]
self.abdomen_z_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_z")]
# Cache arm joint qpos and dof velocity indices for bilateral reach
self.arm_qpos_indices = [
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_x")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_y")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_elbow")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_x")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_y")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_elbow")],
]
self.arm_dof_indices = [
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_x")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_y")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_elbow")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_x")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_y")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_elbow")],
]
# Cache collision geoms for physical impact detection (head, visor, chest, torso)
# All geoms belonging to world body (body 0) are environmental surfaces: floor + all obstacles
self.world_obstacle_geoms = {i for i in range(self.model.ngeom) if self.model.geom_bodyid[i] == 0}
self.body_impact_geoms = {
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "torso_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "chest_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "head_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "visor"),
}
# Cache humanoid geom indices (all geoms belonging to non-world bodies) for ghost rendering
self.humanoid_geoms = [i for i in range(self.model.ngeom) if self.model.geom_bodyid[i] != 0]
# Cache limb geoms for Get-Up multi-phase floor contact righting
self.r_shin_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_shin_geom")
self.l_shin_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_shin_geom")
self.support_hand_geoms = {
self.r_hand_geom_id,
self.l_hand_geom_id,
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_larm_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_larm_geom"),
}
self.support_knee_geoms = {self.r_shin_geom_id, self.l_shin_geom_id}
def _lock_torso_for_reach(self):
"""Bolts root torso, core spine, and legs in neutral athletic stance (Stream 3: Isolated Arm Motor Cortex)."""
self.data.qpos[0] = 0.0
self.data.qpos[1] = 0.0
self.data.qpos[2] = 1.35
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # Upright quaternion
self.data.qvel[0:6] = 0.0 # Zero base linear & angular velocity
# Lock core spine neutral
self.data.qpos[7:10] = 0.0 # abdomen_z, abdomen_y, abdomen_x
# Lock legs in neutral supportive stance
self.data.qpos[10:22] = 0.0 # zero all hip, knee, ankle
self.data.qpos[self.r_knee_qpos] = 0.30
self.data.qpos[self.l_knee_qpos] = 0.30
self.data.qpos[self.r_ankle_qpos] = -0.10
self.data.qpos[self.l_ankle_qpos] = -0.10
self.data.qvel[6:21] = 0.0 # Zero core & lower body velocities (DOFs 6 to 20, preserving arm DOFs 21:27)
def _sample_reach_target(self):
"""
Samples a randomized 3D reach goal via Forward Kinematics (FK) from the active arm's
reachable joint space, mathematically guaranteeing 100.0% of targets are physically reachable (Stream 3).
"""
self.active_reach_hand = np.random.choice(["right", "left"])
# Save current posture
saved_qpos = self.data.qpos.copy()
target_pos = None
# Sample joint angles from the arm's valid forward/lateral reachable envelope
for _ in range(25):
if self.active_reach_hand == "right":
# Right arm envelope: shoulder roll/pitch and elbow flexion
target_sx = float(np.random.uniform(-0.75, 0.40))
target_sy = float(np.random.uniform(-0.60, 0.85))
target_el = float(np.random.uniform(-1.10, 0.45))
self.data.qpos[self.arm_qpos_indices[0]] = target_sx # right_shoulder_x
self.data.qpos[self.arm_qpos_indices[1]] = target_sy # right_shoulder_y
self.data.qpos[self.arm_qpos_indices[2]] = target_el # right_elbow
mujoco.mj_forward(self.model, self.data)
cand_pos = self.data.geom_xpos[self.r_hand_geom_id].copy()
else:
# Left arm envelope: shoulder roll/pitch and elbow flexion
target_sx = float(np.random.uniform(-0.40, 0.75))
target_sy = float(np.random.uniform(-0.60, 0.85))
target_el = float(np.random.uniform(-0.45, 1.10))
self.data.qpos[self.arm_qpos_indices[3]] = target_sx # left_shoulder_x
self.data.qpos[self.arm_qpos_indices[4]] = target_sy # left_shoulder_y
self.data.qpos[self.arm_qpos_indices[5]] = target_el # left_elbow
mujoco.mj_forward(self.model, self.data)
cand_pos = self.data.geom_xpos[self.l_hand_geom_id].copy()
# Restore original pose to measure distance
self.data.qpos[:] = saved_qpos
mujoco.mj_forward(self.model, self.data)
curr_h = self.data.geom_xpos[self.r_hand_geom_id] if self.active_reach_hand == "right" else self.data.geom_xpos[self.l_hand_geom_id]
dist_to_cand = float(np.linalg.norm(curr_h - cand_pos))
target_pos = cand_pos
# Ensure target is at least 0.16m away so the agent must actively articulate its arm
if dist_to_cand >= 0.16:
break
# Always restore state cleanly
self.data.qpos[:] = saved_qpos
mujoco.mj_forward(self.model, self.data)
self.reach_target_pos = target_pos.astype(np.float32)
self.goal_pos = self.reach_target_pos.copy()
def _apply_fall_state(self, pose_name: Optional[str] = None):
"""
Initializes agent in a randomized fall state with 360Β° randomized yaw,
perturbed joints, physics contact settling, and a randomized standing destination goal (Stream 2).
"""
if pose_name is None:
poses = ["prone", "supine", "lateral_left", "lateral_right", "kneeling_crumpled", "seated_slumped", "random_ragdoll"]
pose_name = np.random.choice(poses)
self.fall_state_type = pose_name
self.data.qvel[:] = 0.0
# Random initial floor translation offset
spawn_x = float(np.random.uniform(-0.4, 0.4))
spawn_y = float(np.random.uniform(-0.4, 0.4))
# Base orientation quaternion [w, x, y, z] per canonical pose
if pose_name == "prone":
# Lying on stomach: pitch 90 deg forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.0, 0.7071, 0.0])
self.data.qpos[self.r_knee_qpos] = 0.2
self.data.qpos[self.l_knee_qpos] = 0.2
elif pose_name == "supine":
# Lying on back: pitch 90 deg backward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.0, -0.7071, 0.0])
self.data.qpos[self.r_knee_qpos] = 0.3
self.data.qpos[self.l_knee_qpos] = 0.3
elif pose_name == "lateral_left":
# Lying on left flank: roll 90 deg left
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, -0.7071, 0.0, 0.0])
elif pose_name == "lateral_right":
# Lying on right flank: roll 90 deg right
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.7071, 0.0, 0.0])
elif pose_name == "kneeling_crumpled":
# Crumpled on knees, torso pitched forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.48])
q_base = np.array([0.9238, 0.0, 0.3826, 0.0])
self.data.qpos[self.r_knee_qpos] = 1.7
self.data.qpos[self.l_knee_qpos] = 1.7
self.data.qpos[self.r_hip_qpos] = -1.0
self.data.qpos[self.l_hip_qpos] = -1.0
elif pose_name == "seated_slumped":
# Slumped seated on buttocks, legs forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.35])
q_base = np.array([0.9238, 0.0, -0.3826, 0.0])
self.data.qpos[self.r_hip_qpos] = -1.4
self.data.qpos[self.l_hip_qpos] = -1.4
self.data.qpos[self.r_knee_qpos] = 0.2
self.data.qpos[self.l_knee_qpos] = 0.2
else:
# Random ragdoll drop: drop from 0.45m with completely randomized 3D orientation
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, float(np.random.uniform(0.40, 0.60))])
rand_quat = np.random.uniform(-1.0, 1.0, size=4)
q_base = rand_quat / np.linalg.norm(rand_quat)
# Apply randomized 360Β° yaw rotation to base pose quaternion
yaw = float(np.random.uniform(-np.pi, np.pi))
q_yaw = np.array([np.cos(yaw / 2.0), 0.0, 0.0, np.sin(yaw / 2.0)], dtype=np.float64)
# Quaternion multiplication: q_yaw * q_base
w1, x1, y1, z1 = q_yaw
w2, x2, y2, z2 = q_base
q_rot = np.array([
w1*w2 - x1*x2 - y1*y2 - z1*z2,
w1*x2 + x1*w2 + y1*z2 - z1*y2,
w1*y2 - x1*z2 + y1*w2 + z1*x2,
w1*z2 + x1*y2 - y1*x2 + z1*w2,
], dtype=np.float64)
self.data.qpos[3:7] = q_rot / np.linalg.norm(q_rot)
# Stochastic joint perturbation across all 21 joint degrees of freedom
self.data.qpos[7:] += np.random.uniform(-0.20, 0.20, size=self.n_actuators)
# Settle contact forces smoothly (12 passive sub-steps with zero actuation)
mujoco.mj_forward(self.model, self.data)
for _ in range(12):
self.data.ctrl[:] = 0.0
mujoco.mj_step(self.model, self.data)
self.data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self.data)
# Sample a randomized standing destination goal beacon (0.6m - 1.8m away at random heading)
stand_dist = float(np.random.uniform(0.6, 1.8))
stand_heading = float(np.random.uniform(-np.pi, np.pi))
cur_torso = self.data.xpos[self.torso_body_id]
gx = float(cur_torso[0] + stand_dist * np.cos(stand_heading))
gy = float(cur_torso[1] + stand_dist * np.sin(stand_heading))
gz = 1.168 # Grounded standing torso height equilibrium
self.getup_goal_pos = np.array([gx, gy, gz], dtype=np.float32)
self.goal_pos = self.getup_goal_pos.copy()
def _sample_new_goal(self, is_initial: bool = False):
"""Samples a new random navigation target within the 50m x 50m flat arena."""
torso_pos = self.data.xpos[self.torso_body_id]
if is_initial:
# Initial target: distance 5.0m - 7.5m in forward cone (-60 deg to +60 deg)
dist = float(np.random.uniform(5.0, 7.5))
angle = float(np.random.uniform(-np.pi / 3, np.pi / 3))
self.goal_transition_steps = 0 # Spawn directly facing target: immediate approach required
else:
# Subsequent target: distance 5.5m - 9.0m in full 360 deg omnidirectional bearing
dist = float(np.random.uniform(5.5, 9.0))
angle = float(np.random.uniform(-np.pi, np.pi))
self.goal_transition_steps = self.transition_leeway_window # ~1.85s direction change leeway
gx = float(np.clip(torso_pos[0] + dist * np.cos(angle), -22.0, 22.0))
gy = float(np.clip(torso_pos[1] + dist * np.sin(angle), -22.0, 22.0))
# Ensure newly sampled target is at least (arrival_radius + 1.5m = 3.5m) away from current torso
for _ in range(20):
if np.linalg.norm(np.array([gx - torso_pos[0], gy - torso_pos[1]])) >= (self.goal_arrival_radius + 1.5):
break
center_angle = float(np.arctan2(-torso_pos[1], -torso_pos[0]))
angle = float(center_angle + np.random.uniform(-np.pi / 4, np.pi / 4))
dist = float(np.random.uniform(5.5, 8.5))
gx = float(np.clip(torso_pos[0] + dist * np.cos(angle), -22.0, 22.0))
gy = float(np.clip(torso_pos[1] + dist * np.sin(angle), -22.0, 22.0))
gz = 1.25 # Torso/chest level
self.goal_pos = np.array([gx, gy, gz], dtype=np.float32)
self.min_dist_to_goal = float(np.linalg.norm(self.goal_pos - torso_pos))
self.prev_dist_to_goal = float(self.min_dist_to_goal)
self.steps_since_progress = 0
self.steps_not_approaching = 0
self.negative_vx_steps = 0
@property
def observation_dim(self) -> int:
"""
Total observation vector length:
- Torso height z: 1
- Torso orientation (forward & up vectors): 6
- Torso linear velocity (local/global): 3
- Torso angular velocity: 3
- 21 joint angles: 21
- 21 joint velocities: 21
- 2 foot contact forces: 2
- 512 solid-state LiDAR range rays (16x32 grid): 512
- 3 ego-centric relative waypoint vector (Ξx, Ξy, Ξz local): 3
- 3 metabolic stats (stamina, health, hunger): 3
Total: 1 + 6 + 3 + 3 + 21 + 21 + 2 + 512 + 3 + 3 = 575
"""
return 575
@property
def action_dim(self) -> int:
return self.n_actuators # 21
def reset(self, seed: Optional[int] = None, spawn_x: float = 0.0, initial_vx: float = 0.0) -> np.ndarray:
if seed is not None:
np.random.seed(seed)
mujoco.mj_resetData(self.model, self.data)
# Initial root pose: at X=spawn_x, Y=0, Z=1.33 (athletic bent-knee stance)
self.data.qpos[0] = spawn_x # X
self.data.qpos[1] = 0.0 # Y
self.data.qpos[2] = 1.33 # Z
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # Quaternion (unit)
# Staggered lead leg initialization: alternates 50% left lead, 50% right lead
# This breaks single-leg asymmetric gait traps and forces both thighs to cycle
lead_leg = np.random.choice(["right", "left"])
if lead_leg == "right":
self.data.qpos[self.r_hip_qpos] = -0.35 # Right thigh poised forward
self.data.qpos[self.r_knee_qpos] = 0.45
self.data.qpos[self.r_ankle_qpos] = -0.15
self.data.qpos[self.l_hip_qpos] = 0.05 # Left leg supporting in stance
self.data.qpos[self.l_knee_qpos] = 0.25
self.data.qpos[self.l_ankle_qpos] = -0.05
else:
self.data.qpos[self.l_hip_qpos] = -0.35 # Left thigh poised forward
self.data.qpos[self.l_knee_qpos] = 0.45
self.data.qpos[self.l_ankle_qpos] = -0.15
self.data.qpos[self.r_hip_qpos] = 0.05 # Right leg supporting in stance
self.data.qpos[self.r_knee_qpos] = 0.25
self.data.qpos[self.r_ankle_qpos] = -0.05
# Add slight stochastic noise to joint angles for exploration diversity
joint_noise = np.random.uniform(-0.03, 0.03, size=self.n_actuators)
self.data.qpos[7:] += joint_noise
# Zero velocities with tiny noise (or inject forward running momentum if pioneer spawn)
self.data.qvel[:] = np.random.uniform(-0.02, 0.02, size=self.n_qvel)
if initial_vx > 0.0:
self.data.qvel[0] = initial_vx + np.random.uniform(-0.1, 0.1)
# Reset internal stats
self.stamina = 1.0
self.health = 1.0
self.hunger = 0.0
self.milestones_reached = set()
self.waypoints_reached = set()
self.current_wp_idx = 0
if self.stream_mode == "stand":
# Stream 1: Quiet Standing Balance (Canonical Reference Equilibrium Stance)
self.data.qpos[0] = 0.0
self.data.qpos[1] = 0.0
self.data.qpos[2] = 1.168
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
self.data.qvel[:] = 0.0
self.data.qpos[7:7 + self.n_actuators] = STAND_REF_QPOS.copy()
# Add slight stochastic exploration noise (+/- 0.015 rad ~ 0.8 deg) so agent learns to balance
self.data.qpos[7:7 + self.n_actuators] += np.random.uniform(-0.015, 0.015, size=self.n_actuators)
# Settle foot ground contacts smoothly (10 passive sub-steps with zero actuation)
mujoco.mj_forward(self.model, self.data)
for _ in range(10):
self.data.ctrl[:] = 0.0
mujoco.mj_step(self.model, self.data)
self.data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self.data)
self.goal_pos = np.array([0.0, 0.0, 1.168], dtype=np.float32)
self.stand_still_steps = 0
self.stand_drift_steps = 0
self.stand_one_leg_steps = 0
self.goals_completed = 0
elif self.stream_mode == "getup":
# Stream 2: Fall Recovery (Randomized 360Β° poses & standing destination goal)
self._apply_fall_state()
self.prev_z = float(self.data.xpos[self.torso_body_id][2])
self.prev_head_z = float(self.data.xpos[self.head_body_id][2])
self.consecutive_upright_steps = 0
self.getup_stagnant_steps = 0
self.max_elevation_achieved = self.prev_z
self.max_head_elevation = self.prev_head_z
self.getup_success_count = 0
self.goals_completed = 0
elif self.stream_mode == "reach":
# Stream 2: Bilateral Hand Reaching (Bolted Torso with neutral resting arm configuration)
self._lock_torso_for_reach()
# Initialize both arms down at sides (natural resting posture at hip)
self.data.qpos[self.arm_qpos_indices[0]] = 0.65 + float(np.random.uniform(-0.03, 0.03)) # right_shoulder_x (down at side)
self.data.qpos[self.arm_qpos_indices[1]] = float(np.random.uniform(-0.03, 0.03)) # right_shoulder_y
self.data.qpos[self.arm_qpos_indices[2]] = float(np.random.uniform(-0.03, 0.03)) # right_elbow
self.data.qpos[self.arm_qpos_indices[3]] = -0.65 + float(np.random.uniform(-0.03, 0.03)) # left_shoulder_x (down at side)
self.data.qpos[self.arm_qpos_indices[4]] = float(np.random.uniform(-0.03, 0.03)) # left_shoulder_y
self.data.qpos[self.arm_qpos_indices[5]] = float(np.random.uniform(-0.03, 0.03)) # left_elbow
self.data.qvel[self.arm_dof_indices] = 0.0
mujoco.mj_forward(self.model, self.data)
self._sample_reach_target()
is_right_target = (self.active_reach_hand == "right")
active_h = self.data.geom_xpos[self.r_hand_geom_id] if is_right_target else self.data.geom_xpos[self.l_hand_geom_id]
self.prev_reach_dist = float(np.linalg.norm(active_h - self.reach_target_pos))
self.min_reach_dist = self.prev_reach_dist
self.reach_touches = 0
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
self.goals_completed = 0
elif self.arena_mode:
self.goals_completed = 0
self._sample_new_goal(is_initial=True)
else:
# Fast-forward waypoint index if spawning downfield
for idx, (wp_pos, _, _) in enumerate(COURSE_WAYPOINTS):
if wp_pos[0] > spawn_x + 0.3:
self.current_wp_idx = idx
break
self.spawn_x = float(spawn_x)
self.furthest_x = float(spawn_x)
self.current_step = 0
self.steps_since_progress = 0
self.last_progress_x = float(spawn_x)
self.negative_vx_steps = 0
self.steps_not_approaching = 0
self.r_foot_contact_steps = 0
self.l_foot_contact_steps = 0
mujoco.mj_forward(self.model, self.data)
return self._get_obs()
def _get_obs(self) -> np.ndarray:
# 1. Torso state
torso_xpos = self.data.xpos[self.torso_body_id] # (3,)
torso_xmat = self.data.xmat[self.torso_body_id].reshape(3, 3) # (3, 3)
forward_vec = torso_xmat[:, 0] # Local X axis in world frame
up_vec = torso_xmat[:, 2] # Local Z axis in world frame
torso_cvel = self.data.cvel[self.torso_body_id] # 6D: [ang_vel(3), lin_vel(3)]
ang_vel = torso_cvel[:3]
lin_vel = torso_cvel[3:]
# 2. Joint states (21 angles, 21 velocities)
joint_qpos = self.data.qpos[7:7 + self.n_actuators]
joint_qvel = self.data.qvel[6:6 + self.n_actuators]
# 3. Touch sensors & Task-Specific Contact Indicators
touch_forces = np.zeros(2, dtype=np.float32)
if self.stream_mode == "reach":
# For isolated Reach, legs are bolted stationary. Repurpose touch_forces as explicit Active Hand Indicator:
# [1.0, 0.0] -> Right Hand Active | [0.0, 1.0] -> Left Hand Active
if self.active_reach_hand == "right":
touch_forces[0] = 1.0
touch_forces[1] = 0.0
else:
touch_forces[0] = 0.0
touch_forces[1] = 1.0
else:
if self.model.nsensor >= 2:
touch_forces[0] = np.clip(self.data.sensordata[0] / 100.0, 0.0, 5.0) # Right foot
touch_forces[1] = np.clip(self.data.sensordata[1] / 100.0, 0.0, 5.0) # Left foot
# 4. True Solid-State Angular LiDAR (512 rays from visor/eyes)
lidar_obs = self._compute_lidar_vision()
# 5. Dynamic Goal / Progressive Waypoint (Ego-centric local relative vector)
if self.stream_mode == "stand":
rel_world = self.goal_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
elif self.stream_mode == "getup":
rel_world = self.getup_goal_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 5.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
elif self.stream_mode == "reach":
# Direct closed-loop motor servoing: vector from ACTIVE HAND directly to the reach target orb!
active_h = self.data.geom_xpos[self.r_hand_geom_id] if (self.active_reach_hand == "right") else self.data.geom_xpos[self.l_hand_geom_id]
rel_hand_world = self.reach_target_pos - active_h
rel_hand_local = torso_xmat.T @ rel_hand_world
waypoint_obs = np.array([
np.clip(rel_hand_local[0] / 0.6, -1.0, 1.0),
np.clip(rel_hand_local[1] / 0.6, -1.0, 1.0),
np.clip(rel_hand_local[2] / 0.6, -1.0, 1.0),
], dtype=np.float32)
elif self.arena_mode:
target_pos = self.goal_pos
rel_world = target_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
else:
target_pos = COURSE_WAYPOINTS[self.current_wp_idx][0]
rel_world = target_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
# 6. Metabolic state (stamina, health, hunger/task-channel)
if self.stream_mode == "reach":
hand_indicator = 1.0 if (self.active_reach_hand == "right") else -1.0
metabolic_obs = np.array([self.stamina, self.health, hand_indicator], dtype=np.float32)
else:
metabolic_obs = np.array([self.stamina, self.health, self.hunger], dtype=np.float32)
obs = np.concatenate([
np.array([torso_xpos[2]], dtype=np.float32), # Torso Z (1)
forward_vec.astype(np.float32), # 3
up_vec.astype(np.float32), # 3
lin_vel.astype(np.float32), # 3
ang_vel.astype(np.float32), # 3
joint_qpos.astype(np.float32), # 21
joint_qvel.astype(np.float32), # 21
touch_forces, # 2
lidar_obs, # 512
waypoint_obs, # 3
metabolic_obs, # 3
])
return np.nan_to_num(obs, nan=0.0, posinf=50.0, neginf=-50.0)
def _compute_lidar_vision(self) -> np.ndarray:
"""
True Solid-State Angular LiDAR scan from the robot's visor/eyes.
Casts 128 rays (8 vertical elevation x 16 horizontal azimuth) using MuJoCo's native C raycaster.
Returns normalized range vector (128,) in [0.0, 1.0].
"""
head_pos = self.data.xpos[self.head_body_id]
head_mat = self.data.xmat[self.head_body_id].reshape(3, 3)
# Ray origin: 0.15m forward from head center through visor
ray_origin = np.ascontiguousarray(head_pos + head_mat @ self._visor_offset, dtype=np.float64)
# Transform all 128 precomputed local vectors to world coordinates
# (128, 3) @ (3, 3).T -> (128, 3)
world_vectors = self.lidar_local_vectors @ head_mat.T
normalized_ranges = np.zeros(self.num_lidar_rays, dtype=np.float32)
raw_grid = np.zeros((self.num_lidar_v, self.num_lidar_h), dtype=np.float32)
ray_idx = 0
for row in range(self.num_lidar_v):
for col in range(self.num_lidar_h):
vec = np.ascontiguousarray(world_vectors[ray_idx], dtype=np.float64)
dist = mujoco.mj_ray(self.model, self.data, ray_origin, vec, None, 1, -1, self._geomid_buf)
if dist >= 0:
gid = int(self._geomid_buf[0])
if gid in self.humanoid_geoms:
actual_dist = self.lidar_max_range
else:
actual_dist = min(dist, self.lidar_max_range)
else:
actual_dist = self.lidar_max_range
raw_grid[row, col] = actual_dist
# Zero-centered range: -0.5 (impact / 0m) to +0.5 (clear / 10m), eliminates DC bias
normalized_ranges[ray_idx] = (actual_dist / self.lidar_max_range) - 0.5
ray_idx += 1
self.current_lidar_grid = raw_grid
return normalized_ranges
def _compute_pseudo_vision(self, agent_x: float, agent_y: float, agent_z: float) -> np.ndarray:
"""Backwards compatibility fallback for 1D pseudo-vision."""
return self._compute_lidar_vision()[:16]
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
self.current_step += 1
# Clip action to [-1.0, 1.0]
action = np.clip(action, -1.0, 1.0)
self.data.ctrl[:] = action
if self.stream_mode == "reach":
# For isolated arm reaching, freeze motor commands to core and legs
self.data.ctrl[0:15] = 0.0
# Substep physics
for _ in range(self.frame_skip):
mujoco.mj_step(self.model, self.data)
if self.stream_mode == "reach":
self._lock_torso_for_reach()
# Extract state measurements
torso_xpos = self.data.xpos[self.torso_body_id]
x, y, z = torso_xpos[0], torso_xpos[1], torso_xpos[2]
torso_cvel = self.data.cvel[self.torso_body_id]
vx = torso_cvel[3] # Forward linear velocity
vy = torso_cvel[4]
torso_xmat = self.data.xmat[self.torso_body_id].reshape(3, 3)
up_z = torso_xmat[2, 2] # Cosine of torso vertical tilt
# Update metabolic stats
torque_cost = np.sum(np.square(action))
self.stamina = np.clip(self.stamina - (torque_cost * 0.0003) + 0.0005, 0.0, 1.0)
self.hunger = np.clip(self.hunger + 0.0001, 0.0, 1.0)
# Multi-limb ground contact detection for floor righting and locomotion
contact_r = False
contact_l = False
contact_hand = False
contact_knee = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
g1, g2 = c.geom1, c.geom2
if (g1 == self.r_foot_geom_id and g2 in self.world_obstacle_geoms) or (g2 == self.r_foot_geom_id and g1 in self.world_obstacle_geoms):
contact_r = True
if (g1 == self.l_foot_geom_id and g2 in self.world_obstacle_geoms) or (g2 == self.l_foot_geom_id and g1 in self.world_obstacle_geoms):
contact_l = True
if (g1 in self.support_hand_geoms and g2 in self.world_obstacle_geoms) or (g2 in self.support_hand_geoms and g1 in self.world_obstacle_geoms):
contact_hand = True
if (g1 in self.support_knee_geoms and g2 in self.world_obstacle_geoms) or (g2 in self.support_knee_geoms and g1 in self.world_obstacle_geoms):
contact_knee = True
touch_r = contact_r or (bool(self.data.sensordata[0] > 5.0) if self.model.nsensor >= 1 else False)
touch_l = contact_l or (bool(self.data.sensordata[1] > 5.0) if self.model.nsensor >= 2 else False)
if touch_r:
self.r_foot_contact_steps += 1
if touch_l:
self.l_foot_contact_steps += 1
upright_scale = float(np.clip((z - 0.45) / 0.40, 0.0, 1.0)) # 0 at z<=0.45m, 1.0 at z>=0.85m
in_leeway = False
torso_speed = float(np.linalg.norm(torso_cvel[3:]))
# ββ STREAM 1: QUIET STANDING STILLNESS (REFERENCE POSE SIMILARITY) ββββββββ
if self.stream_mode == "stand":
head_z = float(self.data.xpos[self.head_body_id][2])
dist_from_origin = float(np.linalg.norm([x, y]))
# 1. Joint Configuration Similarity:
# Gaussian kernel comparing 21 actuated joints against canonical STAND_REF_QPOS
current_qpos = self.data.qpos[7:7 + self.n_actuators]
diff_q = current_qpos - STAND_REF_QPOS
weighted_q_mse = float(np.sum(STAND_JOINT_WEIGHTS * (diff_q ** 2)) / np.sum(STAND_JOINT_WEIGHTS))
r_pose = float(np.exp(-3.0 * weighted_q_mse))
# 2. Torso Center-of-Mass Height Similarity (centered at z* = 1.168m):
r_height = float(np.exp(-30.0 * ((z - 1.168) ** 2)))
# 3. Torso Upright Orientation Similarity:
r_upright = float(np.exp(-12.0 * (max(0.0, 1.0 - up_z) ** 2)))
# 4. Multi-Rigid-Body Stillness Similarity (root linear, root angular, joint vel):
lin_vel_sq = float(np.sum(torso_cvel[3:] ** 2))
ang_vel_sq = float(np.sum(torso_cvel[:3] ** 2))
joint_vel_sq = float(np.mean(self.data.qvel[6:6 + self.n_actuators] ** 2))
r_stillness = float(np.exp(-4.0 * (lin_vel_sq + 0.5 * ang_vel_sq + 0.1 * joint_vel_sq)))
# 5. Foot Ground Contact Multiplier:
both_feet_down = bool(touch_r and touch_l)
one_foot_down = bool(touch_r or touch_l)
if both_feet_down:
m_feet = 1.0
elif one_foot_down:
m_feet = 0.45
else:
m_feet = 0.05
# 6. Action Control / Effort Penalty:
r_ctrl = -0.002 * float(np.sum(action ** 2))
# Bounded Gaussian Similarity Reward (Peak = 14.0/step at quiet standing equilibrium)
total_reward = m_feet * (4.0 * r_pose + 3.0 * r_height + 3.0 * r_upright + 4.0 * r_stillness) + r_ctrl
# Track stillness success criteria:
is_stand_still = bool(
r_pose > 0.85 and
r_height > 0.88 and
r_upright > 0.90 and
r_stillness > 0.80 and
both_feet_down
)
if is_stand_still:
self.stand_still_steps += 1
else:
self.stand_still_steps = max(0, self.stand_still_steps - 1)
# Check collision with environmental obstacles/ground using body geoms (head, chest, torso, visor)
is_obstacle_or_ground_impact = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
if (c.geom1 in self.body_impact_geoms and c.geom2 in self.world_obstacle_geoms) or \
(c.geom2 in self.body_impact_geoms and c.geom1 in self.world_obstacle_geoms):
is_obstacle_or_ground_impact = True
break
terminated = False
# Clean early termination only when genuinely fallen or severely displaced:
if z < 0.85 or up_z < 0.55 or is_obstacle_or_ground_impact:
terminated = True
total_reward -= 15.0
elif dist_from_origin > 1.20:
terminated = True
total_reward -= 10.0
# Mastery time limit: 384 steps (~9.6s)
truncated = (self.current_step >= 384)
if truncated and (not terminated) and z > 1.10 and both_feet_down:
total_reward += 50.0 # Mastery bonus for surviving 384 steps upright on both feet
# Detailed metrics for logging & dashboard
torso_roll = abs(float(torso_xmat[2, 1]))
torso_pitch = abs(float(torso_xmat[2, 0]))
spine_lateral = abs(float(self.data.qpos[self.abdomen_x_qpos]))
spine_sagittal = abs(float(self.data.qpos[self.abdomen_y_qpos]))
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": "Quiet Stand Stillness",
"dist_to_wp": float(dist_from_origin), "waypoints_completed": int(self.stand_still_steps > 150),
"goals_completed": int(self.stand_still_steps > 150),
"dist_to_goal": float(dist_from_origin),
"goal_arrival_radius": 1.20,
"goal_pos": [0.0, 0.0, 1.168],
"target_local_pos": [0.0, 0.0, 0.0],
"gaze_alignment": 1.0,
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "stand", "active_reach_hand": "none", "reach_touches": 0,
"getup_success": False,
"both_feet_down": bool(both_feet_down),
"one_foot_down": bool(one_foot_down),
"stand_still_score": float(is_stand_still),
"torso_roll": float(torso_roll),
"torso_pitch": float(torso_pitch),
"spine_lateral": float(spine_lateral),
"spine_sagittal": float(spine_sagittal),
"r_pose": float(r_pose),
"r_height": float(r_height),
"r_upright": float(r_upright),
"r_stillness": float(r_stillness),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# ββ STREAM 2: FALL RECOVERY (GET-UP) ββββββββββββββββββββββββββββββββ
elif self.stream_mode == "getup":
head_z = float(self.data.xpos[self.head_body_id][2])
delta_head_z = head_z - self.prev_head_z
self.prev_head_z = head_z
r_foot_z = float(self.data.xpos[self.r_foot_id][2])
l_foot_z = float(self.data.xpos[self.l_foot_id][2])
max_foot_z = max(r_foot_z, l_foot_z)
# 1. REWARD FOR HEAD & TORSO ELEVATION (Rises vertically, head leads above torso):
reward_head_elevation = 16.0 * delta_head_z
reward_torso_elevation = 12.0 * max(0.0, z - self.prev_z)
reward_head_height = 5.0 * max(0.0, head_z - 0.40)
reward_head_above_torso = 3.5 * max(0.0, head_z - z)
reward_upright = 3.0 * max(0.0, up_z - 0.30)
# 2. MULTI-PHASE GROUND CONTACT SUPPORT (Hands, Knees, Feet):
both_feet_down = touch_r and touch_l
one_foot_down = touch_r or touch_l
# Floor bracing support (head_z < 1.15m): reward hands pushing floor & knees tucked
reward_hand_brace = 4.0 if (contact_hand and head_z < 1.15) else 0.0
reward_knee_brace = 3.0 if (contact_knee and head_z < 1.15) else 0.0
# Standing stage (head_z >= 1.05m): reward planting feet
if both_feet_down:
reward_feet_ground = 6.0
elif one_foot_down:
reward_feet_ground = 2.0
elif head_z >= 1.20:
reward_feet_ground = -1.0 # Only penalize missing foot contact once upright
else:
reward_feet_ground = 0.0 # Zero penalty while floor-bound rolling or pushing up!
# Floor proximity for feet during ground phase
reward_feet_down = 2.0 * max(0.0, 0.30 - max_foot_z) if head_z < 1.15 else 0.0
# 3. MITIGATE LEG AIR-FLAILING:
# Only penalize feet higher than head IF head is actually on the ground and feet are violently kicked up
feet_above_head = max(0.0, max_foot_z - max(0.40, head_z))
penalty_shooting_legs = -4.0 * feet_above_head
# Air-flail penalty only when feet are high in air with zero contact
penalty_air_flail = -2.0 * max(0.0, max_foot_z - 0.50) if (head_z < 1.0 and not (contact_hand or contact_knee)) else 0.0
# 4. Energy cost & Arm Windmilling Suppression
arm_flail_penalty = -0.001 * float(np.sum(np.square(action[15:21])))
reward_ctrl = -0.001 * torque_cost + arm_flail_penalty
dist_to_stand_goal = float(np.linalg.norm(self.getup_goal_pos[:2] - np.array([x, y])))
reward_goal_heading = 2.0 * max(0.0, 1.0 - dist_to_stand_goal / 2.0) if head_z > 1.30 else 0.0
total_reward = (
reward_head_elevation +
reward_torso_elevation +
reward_head_height +
reward_head_above_torso +
reward_upright +
reward_hand_brace +
reward_knee_brace +
reward_feet_ground +
reward_feet_down +
penalty_shooting_legs +
penalty_air_flail +
reward_ctrl +
reward_goal_heading
)
self.prev_z = z
# Track peak head elevation for stagnation monitoring
if head_z > self.max_head_elevation:
self.max_head_elevation = head_z
self.getup_stagnant_steps = 0
else:
self.getup_stagnant_steps += 1
# Upright standing with feet planted verification:
# Requires head up (> 1.35m), torso upright (z > 1.05m, up_z > 0.70), and foot contact
is_upright = (head_z > 1.35) and (z > 1.05) and (up_z > 0.70) and (one_foot_down or (head_z > 1.45 and z > 1.10))
if is_upright:
self.consecutive_upright_steps += 1
total_reward += 3.0
if both_feet_down:
total_reward += 3.0 # Extra jackpot bonus when BOTH feet are planted!
else:
self.consecutive_upright_steps = 0
# Lenient Ground Collapse Stagnation & Rise Stall Deaths:
# Tolerates 140 steps (~3.5s) struggling on ground without head elevation
# Tolerates 180 steps (~4.5s) coordinating knee-tuck and feet planting to reach upright stance
is_ground_stuck = (head_z < 0.50 and self.getup_stagnant_steps > 140)
is_stall_stuck = (head_z >= 0.50 and self.consecutive_upright_steps == 0 and self.getup_stagnant_steps > 180)
terminated = False
if self.consecutive_upright_steps >= 12:
total_reward += 100.0 # Upright Recovery Jackpot!
self.getup_success_count += 1
terminated = True
elif is_ground_stuck or is_stall_stuck:
terminated = True
total_reward -= 10.0 # Timeout stagnation death
truncated = (self.current_step >= 280)
target_pos = self.getup_goal_pos
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": f"Get-Up ({self.fall_state_type})",
"dist_to_wp": float(max(0.0, 1.55 - head_z)),
"waypoints_completed": int(self.consecutive_upright_steps >= 12),
"goals_completed": int(self.consecutive_upright_steps >= 12),
"dist_to_goal": float(dist_to_stand_goal),
"goal_arrival_radius": 0.50,
"goal_pos": self.getup_goal_pos.tolist(),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(max(0.0, up_z)),
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "getup", "active_reach_hand": "none", "reach_touches": 0,
"getup_success": bool(self.consecutive_upright_steps >= 12),
"head_z": float(head_z),
"both_feet_down": bool(both_feet_down),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# ββ STREAM 2: BILATERAL HAND REACHING (BOLTED BODY) ββββββββββββββββ
elif self.stream_mode == "reach":
# Torso, core, and legs are bolted in neutral athletic upright stance.
# Control is isolated strictly to the 6 arm actuators reaching for 3D orbs.
r_h = self.data.geom_xpos[self.r_hand_geom_id]
l_h = self.data.geom_xpos[self.l_hand_geom_id]
dist_r = float(np.linalg.norm(r_h - self.reach_target_pos))
dist_l = float(np.linalg.norm(l_h - self.reach_target_pos))
# The designated active hand is the hand closer to the target ball
is_right_target = (self.active_reach_hand == "right")
if is_right_target:
active_h = r_h
wrong_h = l_h
dist_to_target = dist_r
dist_inactive = dist_l
active_body_id = self.r_arm_body_id
wrong_body_id = self.l_arm_body_id
active_act = action[15:18]
wrong_act = action[18:21]
active_qvel = self.data.qvel[self.arm_dof_indices[0:3]]
wrong_qvel = self.data.qvel[self.arm_dof_indices[3:6]]
wrong_qpos = self.data.qpos[self.arm_qpos_indices[3:6]]
wrong_rest_qpos = np.array([-0.65, 0.0, 0.0], dtype=np.float32)
else:
active_h = l_h
wrong_h = r_h
dist_to_target = dist_l
dist_inactive = dist_r
active_body_id = self.l_arm_body_id
wrong_body_id = self.r_arm_body_id
active_act = action[18:21]
wrong_act = action[15:18]
active_qvel = self.data.qvel[self.arm_dof_indices[3:6]]
wrong_qvel = self.data.qvel[self.arm_dof_indices[0:3]]
wrong_qpos = self.data.qpos[self.arm_qpos_indices[0:3]]
wrong_rest_qpos = np.array([0.65, 0.0, 0.0], dtype=np.float32)
# 1. Main reach distance progress and proximity reward
dist_prog = self.prev_reach_dist - dist_to_target
reward_reach_prog = 25.0 * dist_prog
reward_proximity = 6.0 * np.exp(-4.0 * dist_to_target)
# 2. MASSIVE SPEED EFFICIENCY REWARDS:
# a) Direct closing velocity vector along the unit ray connecting active hand to target orb
delta_active_target = self.reach_target_pos - active_h
dist_hand_target = max(1e-4, float(np.linalg.norm(delta_active_target)))
u_reach_target = delta_active_target / dist_hand_target
active_hand_vel = self.data.cvel[active_body_id][3:6]
v_closing = float(np.dot(active_hand_vel, u_reach_target))
# Speed efficiency bonus for thrusting active hand toward target
reward_closing_speed = 4.0 * max(0.0, v_closing)
# b) Mild latency penalty (encourages prompt reaching)
penalty_latency = -0.02 * (float(self.steps_on_reach_target) / 40.0)
# 3. Slight reward for moving the correct hand
active_hand_speed = float(np.linalg.norm(active_hand_vel))
active_qvel_mag = float(np.sum(np.abs(active_qvel)))
reward_correct_motion = 0.15 * min(1.0, active_hand_speed / 0.25) + 0.10 * np.tanh(active_qvel_mag)
# 4. REGULARIZATION ON INACTIVE HAND (stay neutral and relaxed at side)
wrong_hand_speed = float(np.linalg.norm(self.data.cvel[wrong_body_id][3:6]))
penalty_wrong_act = -0.05 * float(np.sum(np.square(wrong_act)))
penalty_wrong_qvel = -0.05 * float(np.sum(np.abs(wrong_qvel)))
penalty_wrong_speed = -0.20 * min(1.0, wrong_hand_speed)
penalty_wrong_raised = -0.50 * max(0.0, float(wrong_h[2]) - 1.28)
penalty_wrong_posture = -0.20 * min(1.0, float(np.linalg.norm(wrong_qpos - wrong_rest_qpos)))
penalty_wrong_lead = -0.50 * max(0.0, dist_to_target - dist_inactive)
penalty_wrong_encroach = -0.50 * max(0.0, 0.35 - dist_inactive)
penalty_wrong_touch = -2.0 if dist_inactive < 0.14 else 0.0
# General control penalty on active arm to encourage efficient articulation
reward_ctrl = -0.001 * float(np.sum(np.square(active_act)))
self.steps_on_reach_target += 1
if dist_to_target < self.min_reach_dist:
self.min_reach_dist = dist_to_target
self.steps_receding_reach = 0
elif dist_to_target > self.min_reach_dist + 0.25:
self.steps_receding_reach += 1
reward_touch = 0.0
speed_efficiency_jackpot = 0.0
if dist_to_target < 0.14:
# 1. Base touch award
base_touch = 40.0
# 2. Speed efficiency jackpot: maximum when acquired in few steps
speed_efficiency_jackpot = max(0.0, 110.0 - float(self.steps_on_reach_target)) * 0.40
# 3. Consecutive hit streak multiplier
touch_streak_bonus = min(30.0, float(self.reach_touches) * 6.0)
reward_touch = base_touch + speed_efficiency_jackpot + touch_streak_bonus
self.reach_touches += 1
self._sample_reach_target()
is_right_target = (self.active_reach_hand == "right")
active_h = self.data.geom_xpos[self.r_hand_geom_id] if is_right_target else self.data.geom_xpos[self.l_hand_geom_id]
dist_to_target = float(np.linalg.norm(active_h - self.reach_target_pos))
# Reset timers for next target!
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
self.min_reach_dist = dist_to_target
self.prev_reach_dist = dist_to_target
total_reward = (
reward_reach_prog +
reward_proximity +
reward_closing_speed +
reward_correct_motion +
reward_touch +
penalty_latency +
penalty_wrong_act +
penalty_wrong_qvel +
penalty_wrong_speed +
penalty_wrong_raised +
penalty_wrong_posture +
penalty_wrong_lead +
penalty_wrong_encroach +
penalty_wrong_touch +
reward_ctrl
)
# Lenient Reaching Time Limits:
# Tolerates 110 steps (~2.75s) per target orb, or 45 steps receding away
is_reach_timeout = (self.steps_on_reach_target > 110)
is_reach_receding = (self.steps_receding_reach > 45)
terminated = False
if is_reach_timeout or is_reach_receding:
terminated = True
total_reward -= 2.0 # Mild timeout penalty
truncated = (self.current_step >= 384)
target_pos = self.reach_target_pos
head_xmat = self.data.xmat[self.head_body_id].reshape(3, 3)
head_fwd = head_xmat[:, 0]
delta_head_target = target_pos - self.data.xpos[self.head_body_id]
cos_gaze = float(np.dot(head_fwd, delta_head_target / max(1e-3, np.linalg.norm(delta_head_target))))
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": f"Reach ({self.active_reach_hand.upper()})",
"dist_to_wp": float(dist_to_target), "waypoints_completed": int(self.reach_touches),
"goals_completed": int(self.reach_touches),
"dist_to_goal": float(dist_to_target),
"goal_arrival_radius": 0.14,
"goal_pos": self.reach_target_pos.tolist(),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(cos_gaze),
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "reach", "active_reach_hand": str(self.active_reach_hand),
"reach_touches": int(self.reach_touches), "getup_success": False, "stand_still_score": 0.0,
"wrong_hand_dist": float(dist_inactive), "wrong_hand_touch": bool(dist_inactive < 0.14),
"wrong_hand_speed": float(wrong_hand_speed),
"wrong_arm_action_penalty": float(penalty_wrong_act),
"wrong_arm_raised_penalty": float(penalty_wrong_raised),
"v_closing": float(v_closing),
"reward_closing_speed": float(reward_closing_speed),
"speed_efficiency_jackpot": float(speed_efficiency_jackpot),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# Target Navigation Vector & Forward Propulsion
if self.arena_mode:
target_pos = self.goal_pos
delta_target = target_pos - torso_xpos
dist_to_target = float(np.linalg.norm(delta_target))
dist_xy = float(np.linalg.norm(delta_target[:2]))
# Check dynamic goal arrival (entering wide arrival radius)
reward_arrival = 0.0
if dist_xy < self.goal_arrival_radius:
self.goals_completed += 1
# Escalating jackpot for chained consecutive goals: 50 for 1st, 85 for 2nd, 120 for 3rd, 155 for 4th...
reward_arrival = 50.0 + 35.0 * (self.goals_completed - 1)
self._sample_new_goal(is_initial=False)
target_pos = self.goal_pos
delta_target = target_pos - torso_xpos
dist_to_target = float(np.linalg.norm(delta_target))
dist_xy = float(np.linalg.norm(delta_target[:2]))
self.prev_dist_to_goal = dist_to_target
dist_delta = 0.0
else:
dist_delta = self.prev_dist_to_goal - dist_to_target
self.prev_dist_to_goal = dist_to_target
# Leeway period active after goal completion to allow decelerating & turning
in_leeway = (self.goal_transition_steps > 0)
if in_leeway:
self.goal_transition_steps -= 1
# While turning during leeway, continuously update closest distance baseline to current position
self.min_dist_to_goal = dist_to_target
# 1. Distance progress potential (scaled with goal streak) & sharp penalty for moving away
streak_mult = 1.0 + 0.35 * min(4, self.goals_completed)
if dist_delta >= 0.0:
reward_dist_progress = 3.0 * streak_mult * dist_delta
else:
# Double penalty for moving further away from target outside of turn leeway
reward_dist_progress = 0.0 if in_leeway else (6.0 * dist_delta)
# 2. Penalty for drifting/retreating beyond the closest distance achieved so far (muted in leeway)
retreat_distance = max(0.0, dist_to_target - self.min_dist_to_goal)
penalty_retreat = 0.0 if in_leeway else (-3.0 * retreat_distance if retreat_distance > 0.10 else 0.0)
unit_target_dir = delta_target / max(0.1, dist_to_target)
vel_towards_target = float(np.dot(torso_cvel[3:], unit_target_dir))
torso_speed = float(np.linalg.norm(torso_cvel[3:]))
# 3. Controlled Forward Movement vs Severe Retrograde Penalty
# Cap forward reward at 1.8 m/s so Olympic sprinting gives no extra reward
v_eff = min(max(0.0, vel_towards_target), 1.8)
reward_forward = 3.0 * v_eff * upright_scale
penalty_retrograde = 0.0 if (vel_towards_target >= 0.0 or in_leeway) else (5.0 * vel_towards_target)
# 4. Heading guidance velocity (steep negative slope when heading backwards outside leeway)
if vel_towards_target >= 0.0:
reward_heading = 1.5 * v_eff * upright_scale
else:
reward_heading = 0.0 if in_leeway else (3.0 * vel_towards_target)
# 5. Speed Regulation Penalty: Penalize reckless overspeeding (> 2.0 m/s) to promote stable locomotion
penalty_overspeed = 0.0
if torso_speed > 2.0:
penalty_overspeed = -3.0 * ((torso_speed - 2.0) ** 2)
# 6. Turn alignment & braking reward during transition leeway: encourage slowing down & rapid yaw alignment to new goal
reward_turn = 0.0
reward_brake = 0.0
if in_leeway:
torso_fwd = torso_xmat[:, 0]
cos_torso_goal = float(np.dot(torso_fwd[:2], unit_target_dir[:2]))
# Reward rotating body to face new target (+0.0 to +1.0)
reward_turn = 1.0 * max(0.0, cos_torso_goal)
# Reward decelerating below 1.6 m/s so the agent can pivot cleanly
if torso_speed < 1.6:
reward_brake = 0.6 * (1.6 - torso_speed)
# 7. Subsequent Goal Urgency: Time-decay pressure once Goal 1 is cleared
penalty_urgency = 0.0
if self.goals_completed >= 1 and not in_leeway:
penalty_urgency = -0.08 * min(3, self.goals_completed)
reward_wp = (
reward_arrival +
reward_dist_progress +
penalty_retreat +
penalty_retrograde +
penalty_overspeed +
reward_turn +
reward_brake +
penalty_urgency
)
reward_drift = 0.0
reward_milestone = 0.0
dist_to_wp = dist_to_target
else:
# Obstacle course mode
wp_pos, wp_name, wp_radius = COURSE_WAYPOINTS[self.current_wp_idx]
target_pos = wp_pos
delta_to_wp = wp_pos - torso_xpos
dist_to_wp = float(np.linalg.norm(delta_to_wp))
dist_xy = float(np.linalg.norm(delta_to_wp[:2]))
reward_forward = 3.5 * vx * upright_scale
if vx < -0.05:
reward_forward += 1.5 * vx
wp_reached = False
if dist_xy < wp_radius:
wp_reached = True
elif torso_xpos[0] >= wp_pos[0] and abs(torso_xpos[1] - wp_pos[1]) < (wp_radius + 0.3):
wp_reached = True
reward_wp = 0.0
if wp_reached and self.current_wp_idx not in self.waypoints_reached:
self.waypoints_reached.add(self.current_wp_idx)
if self.current_wp_idx == len(COURSE_WAYPOINTS) - 1:
reward_wp = 100.0
elif 3 <= self.current_wp_idx <= 14:
reward_wp = 25.0
else:
reward_wp = 15.0
if self.current_wp_idx < len(COURSE_WAYPOINTS) - 1:
self.current_wp_idx += 1
wp_pos, wp_name, wp_radius = COURSE_WAYPOINTS[self.current_wp_idx]
target_pos = wp_pos
delta_to_wp = wp_pos - torso_xpos
dist_to_wp = float(np.linalg.norm(delta_to_wp))
wp_unit_dir = delta_to_wp / max(0.1, dist_to_wp)
vel_towards_wp = float(np.dot(torso_cvel[3:], wp_unit_dir))
reward_heading = 2.0 * max(-0.5, min(vel_towards_wp, 3.5)) * upright_scale
target_y = float(wp_pos[1])
reward_drift = -0.5 * abs(y - target_y)
reward_milestone = 0.0
for threshold, name in STAGE_MILESTONES:
if x >= threshold and threshold not in self.milestones_reached:
self.milestones_reached.add(threshold)
reward_milestone += 50.0
# Head Gaze Alignment Reward: mild reward for facing objective with head/visor
head_xmat = self.data.xmat[self.head_body_id].reshape(3, 3)
head_fwd = head_xmat[:, 0] # Local X axis of head (forward gaze direction)
head_pos = self.data.xpos[self.head_body_id]
delta_head_target = target_pos - head_pos
dist_head_target = float(np.linalg.norm(delta_head_target))
u_head_target = delta_head_target / max(1e-3, dist_head_target)
cos_gaze = float(np.dot(head_fwd, u_head_target))
reward_gaze = 0.15 * max(0.0, cos_gaze)
# 2. Upright posture reward (staying balanced above 0.70m)
reward_upright = 1.0 if z > 0.70 else 0.0
# 3. Upright orientation bonus & active lateral roll stabilization
roll_y = torso_xmat[1, 2]
reward_roll = -2.0 * abs(roll_y)
reward_orient = 0.8 * max(0.0, up_z) + reward_roll
# Advance speed for biomechanical stepping gating
forward_vec = torso_xmat[:, 0]
v_advance = float(np.dot(torso_cvel[3:6], forward_vec)) if self.arena_mode else vx
# 4. Alternating Bipedal Stepping Cadence
reward_gait = 0.0
if z > 0.55 and v_advance > 0.10:
if touch_r ^ touch_l:
reward_gait = 2.0 * min(v_advance, 1.5)
elif touch_r and touch_l:
reward_gait = -0.4 * min(v_advance, 1.5)
# 5. Knee Articulation & Stance Stability
r_knee_q = self.data.qpos[self.r_knee_qpos]
l_knee_q = self.data.qpos[self.l_knee_qpos]
reward_knee = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r:
if 0.25 <= r_knee_q <= 1.1:
reward_knee += 0.3
elif r_knee_q > 1.6:
reward_knee -= 0.5 * (r_knee_q - 1.6)
if not touch_l:
if 0.25 <= l_knee_q <= 1.1:
reward_knee += 0.3
elif l_knee_q > 1.6:
reward_knee -= 0.5 * (l_knee_q - 1.6)
if touch_r and r_knee_q < 0.50:
reward_knee += 0.2
if touch_l and l_knee_q < 0.50:
reward_knee += 0.2
# 6. Swing foot forward velocity bonus
r_foot_vfwd = float(np.dot(self.data.cvel[self.r_foot_id][3:6], forward_vec)) if self.arena_mode else self.data.cvel[self.r_foot_id][3]
l_foot_vfwd = float(np.dot(self.data.cvel[self.l_foot_id][3:6], forward_vec)) if self.arena_mode else self.data.cvel[self.l_foot_id][3]
reward_swing = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r and r_foot_vfwd > 0.15:
reward_swing += 0.3 * min(r_foot_vfwd, 2.0)
if not touch_l and l_foot_vfwd > 0.15:
reward_swing += 0.3 * min(l_foot_vfwd, 2.0)
# 7. Thigh Forward Stride Separation
r_hip_q = self.data.qpos[self.r_hip_qpos]
l_hip_q = self.data.qpos[self.l_hip_qpos]
reward_stride = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r:
stride_diff = l_hip_q - r_hip_q
if 0.1 < stride_diff < 1.3:
reward_stride = 1.0 * stride_diff
elif not touch_l:
stride_diff = r_hip_q - l_hip_q
if 0.1 < stride_diff < 1.3:
reward_stride = 1.0 * stride_diff
# 8. Bilateral Stance Symmetry (Duty-Cycle Equalizer)
reward_symmetry = 0.0
total_contacts = self.r_foot_contact_steps + self.l_foot_contact_steps
if self.current_step > 10 and total_contacts > 6 and z > 0.55:
asym = abs(self.r_foot_contact_steps - self.l_foot_contact_steps) / total_contacts
if asym > 0.20:
reward_symmetry = -2.0 * (asym - 0.20)
# 9. Arm Windmilling Suppression (Penalize flailing arms to force balance from core and legs)
arm_flail_penalty = -0.008 * float(np.sum(np.square(action[15:21])))
# 10. Energy expenditure penalty
reward_ctrl = -0.0015 * torque_cost + arm_flail_penalty
total_reward = (
reward_forward +
reward_heading +
reward_wp +
reward_gaze +
reward_upright +
reward_orient +
reward_gait +
reward_knee +
reward_swing +
reward_stride +
reward_symmetry +
reward_ctrl +
reward_drift +
reward_milestone
)
if self.arena_mode:
rad_dist = float(np.linalg.norm([x, y]))
if rad_dist > self.furthest_x:
self.furthest_x = rad_dist
if in_leeway:
# During turning leeway window: reset & do not accumulate moving-away or stagnation counters
self.steps_not_approaching = 0
self.negative_vx_steps = 0
self.steps_since_progress = 0
self.min_dist_to_goal = dist_to_target
else:
# Active approach check: requires advancing toward target (> 0.15 m/s and closing distance)
is_actively_approaching = (vel_towards_target >= 0.15) and (dist_delta > 0.001)
if is_actively_approaching:
self.steps_not_approaching = 0
else:
self.steps_not_approaching += 1
# Receding tracker: moving away from target
is_receding = (vel_towards_target < -0.05) or (dist_delta < -0.003)
if is_receding:
self.negative_vx_steps += 1
else:
self.negative_vx_steps = max(0, self.negative_vx_steps - 1)
# Record best approach distance
if dist_to_target < self.min_dist_to_goal:
self.min_dist_to_goal = dist_to_target
self.steps_since_progress = 0
else:
self.steps_since_progress += 1
else:
if x > self.furthest_x:
self.furthest_x = x
if x > self.last_progress_x + 0.06:
self.last_progress_x = x
self.steps_since_progress = 0
else:
self.steps_since_progress += 1
if vx < -0.15:
self.negative_vx_steps += 1
else:
self.negative_vx_steps = 0
# ββ TERMINATION & TRUNCATION ββββββββββββββββββββββββββββββββββββββββββ
# Exact preserved death conditions:
# 1. Physical collision: head, visor, chest, or upper torso directly impacts the floor/obstacles
# 2. True collapse: torso center of mass drops below 0.35m or head drops below 0.45m
# 3. Complete inversion: torso tilted > 81 deg from vertical (up_z < 0.15)
# 4. Arena boundary departure (|x| > 24m or |y| > 24m) or course lateral departure (|y| > 3.0m)
# 5. Fast elimination when moving away or not actively approaching target (outside leeway)
terminated = False
is_obstacle_or_ground_impact = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
g1, g2 = c.geom1, c.geom2
if (g1 in self.body_impact_geoms and g2 in self.world_obstacle_geoms) or \
(g2 in self.body_impact_geoms and g1 in self.world_obstacle_geoms):
is_obstacle_or_ground_impact = True
break
head_z = self.data.xpos[self.head_body_id][2]
is_altitude_collapse = (z < 0.35) or (head_z < 0.45)
is_tilt_collapse = (up_z < 0.15)
is_physical_fall = is_obstacle_or_ground_impact or is_altitude_collapse or is_tilt_collapse
if self.arena_mode:
is_boundary_departure = (abs(x) > 24.0 or abs(y) > 24.0)
# FAST KILL TRIGGERS (active outside turning leeway period):
# 1. Backing away kill: moved away from target for > 25 steps (~0.6s)
is_moving_away_kill = (not in_leeway) and (self.negative_vx_steps > 25)
# 2. Retreat distance kill: backed up > 0.40m further than closest approach achieved
is_retreat_kill = (not in_leeway) and (dist_to_target > self.min_dist_to_goal + 0.40) and (self.current_step > 15)
# 3. Not actively approaching kill: dawdling, circling, or stationary for > 80 steps (~2.0s)
is_not_approaching_kill = (not in_leeway) and (self.current_step > 25 and self.steps_not_approaching > 80)
# 4. Stagnation kill: no new record distance achieved for > 85 steps (~2.125s)
is_stagnant_kill = (not in_leeway) and (self.current_step > 25 and self.steps_since_progress > 85)
if is_physical_fall:
terminated = True
total_reward -= 20.0
self.health = max(0.0, self.health - 0.2)
elif is_boundary_departure:
terminated = True
total_reward -= 20.0
elif is_moving_away_kill or is_retreat_kill:
terminated = True
total_reward -= 25.0 # Severe penalty for retreating/moving away from target
elif is_not_approaching_kill or is_stagnant_kill:
terminated = True
total_reward -= 25.0 # Severe penalty for failing to actively approach target
# Dropped Goal Streak Penalty:
# If the agent achieved at least 1 goal, but died before reaching the next goal,
# apply a moderate penalty to encourage multi-goal chaining without negating arrival returns
if terminated and self.goals_completed >= 1:
dropped_goal_penalty = -15.0 * (1.0 + 0.2 * min(3, self.goals_completed - 1))
total_reward += dropped_goal_penalty
else:
is_moving_backwards = (x < self.spawn_x - 0.50) or (self.negative_vx_steps > 60)
is_stagnant = (self.current_step > 50 and self.steps_since_progress > 90)
is_lateral_departure = (abs(y) > 3.0)
if is_physical_fall:
terminated = True
total_reward -= 20.0
self.health = max(0.0, self.health - 0.2)
elif is_moving_backwards:
terminated = True
total_reward -= 15.0
elif is_stagnant:
terminated = True
total_reward -= 10.0
elif is_lateral_departure:
terminated = True
total_reward -= 20.0
if x >= 47.0:
terminated = True
total_reward += 200.0
truncated = (self.current_step >= self.max_steps)
info = {
"x_pos": float(x),
"y_pos": float(y),
"z_pos": float(z),
"vx": float(vx),
"torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x),
"milestones": len(self.milestones_reached),
"waypoint_idx": int(self.current_wp_idx),
"waypoint_name": str(COURSE_WAYPOINTS[self.current_wp_idx][1]) if not self.arena_mode else f"Arena Target #{self.goals_completed + 1}",
"dist_to_wp": float(dist_to_wp),
"waypoints_completed": len(self.waypoints_reached) if not self.arena_mode else int(self.goals_completed),
"goals_completed": int(self.goals_completed),
"dist_to_goal": float(dist_to_wp),
"goal_arrival_radius": float(self.goal_arrival_radius) if self.arena_mode else 1.0,
"goal_pos": self.goal_pos.tolist() if self.arena_mode else list(COURSE_WAYPOINTS[self.current_wp_idx][0]),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(cos_gaze),
"stamina": float(self.stamina),
"health": float(self.health),
"hunger": float(self.hunger),
"step": self.current_step,
"lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": bool(in_leeway) if self.arena_mode else False,
"leeway_remaining": int(self.goal_transition_steps) if self.arena_mode else 0,
"stream_mode": "cruise",
"active_reach_hand": "none",
"reach_touches": 0,
"getup_success": False,
"stand_still_score": 0.0,
}
return self._get_obs(), float(total_reward), terminated, truncated, info
def render_frame(
self,
width: int = 480,
height: int = 320,
camera_name: str = "behind_cam",
ghost_datas: Optional[List[Any]] = None,
ghost_alpha: Optional[float] = None,
) -> np.ndarray:
"""
Renders an RGB numpy image of the current simulation frame.
Optionally injects translucent ghost avatars into mjvScene for parallel multi-agent swarm visualization.
"""
if self._renderer is None or self._renderer.width != width or self._renderer.height != height:
self._renderer = mujoco.Renderer(self.model, height=height, width=width)
self._renderer.update_scene(self.data, camera=camera_name)
# Inject dynamic course waypoint goal markers into mjvScene
scene = self._renderer.scene
identity_mat = np.eye(3).flatten()
if self.stream_mode == "reach":
# Stream 3 Reach: Glowing magenta hand target orb + vertical guide line
rx, ry, rz = float(self.reach_target_pos[0]), float(self.reach_target_pos[1]), float(self.reach_target_pos[2])
if scene.ngeom < scene.maxgeom:
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([rx, ry, rz], dtype=np.float64)
orb_size = np.array([0.09, 0.09, 0.09], dtype=np.float64)
orb_rgba = np.array([0.88, 0.20, 0.95, 0.95], dtype=np.float32) # Radiant Magenta
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
if scene.ngeom < scene.maxgeom:
dst_guide = scene.geoms[scene.ngeom]
guide_pos = np.array([rx, ry, max(0.01, rz / 2.0)], dtype=np.float64)
guide_size = np.array([0.012, max(0.01, rz / 2.0), 0.0], dtype=np.float64)
guide_rgba = np.array([0.88, 0.20, 0.95, 0.35], dtype=np.float32)
mujoco.mjv_initGeom(dst_guide, mujoco.mjtGeom.mjGEOM_CYLINDER, guide_size, guide_pos, identity_mat, guide_rgba)
dst_guide.transparent = 1
scene.ngeom += 1
elif self.stream_mode == "getup":
# Stream 2 Get-Up: Golden upright elevation target disc
torso_xpos = self.data.xpos[self.torso_body_id]
tx, ty = float(torso_xpos[0]), float(torso_xpos[1])
if scene.ngeom < scene.maxgeom:
dst_target = scene.geoms[scene.ngeom]
target_pos = np.array([tx, ty, 1.33], dtype=np.float64)
target_size = np.array([0.25, 0.02, 0.0], dtype=np.float64)
target_rgba = np.array([0.98, 0.75, 0.15, 0.65], dtype=np.float32) # Glowing Amber
mujoco.mjv_initGeom(dst_target, mujoco.mjtGeom.mjGEOM_CYLINDER, target_size, target_pos, identity_mat, target_rgba)
dst_target.transparent = 1
scene.ngeom += 1
elif self.stream_mode == "stand":
# Stream 1 Stand: Cyan stillness equilibrium ground pad
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([0.0, 0.0, 0.015], dtype=np.float64)
ring_size = np.array([0.35, 0.015, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.50], dtype=np.float32) # Cyan Pad
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
elif self.arena_mode:
# ACTIVE ARENA GOAL: Glowing floating target orb + vertical laser beacon column + ground target ring
gx, gy, gz = float(self.goal_pos[0]), float(self.goal_pos[1]), float(self.goal_pos[2])
# 1. Floating target orb (sphere)
if scene.ngeom < scene.maxgeom:
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([gx, gy, gz + 0.35], dtype=np.float64)
orb_size = np.array([0.22, 0.22, 0.22], dtype=np.float64)
orb_rgba = np.array([0.98, 0.75, 0.15, 0.90], dtype=np.float32) # Vibrant Gold
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
# 2. Vertical laser beacon column (cylinder)
if scene.ngeom < scene.maxgeom:
dst_beam = scene.geoms[scene.ngeom]
beam_pos = np.array([gx, gy, 1.25], dtype=np.float64)
beam_size = np.array([0.05, 1.25, 0.0], dtype=np.float64)
beam_rgba = np.array([0.22, 0.74, 0.97, 0.55], dtype=np.float32) # Glowing Cyan
mujoco.mjv_initGeom(dst_beam, mujoco.mjtGeom.mjGEOM_CYLINDER, beam_size, beam_pos, identity_mat, beam_rgba)
dst_beam.transparent = 1
scene.ngeom += 1
# 3. Ground target arrival ring
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([gx, gy, 0.02], dtype=np.float64)
ring_size = np.array([self.goal_arrival_radius, 0.02, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.45], dtype=np.float32)
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
else:
for wp_idx, (wp_coord, wp_name, wp_radius) in enumerate(COURSE_WAYPOINTS):
if scene.ngeom + 3 >= scene.maxgeom:
break
wx = float(wp_coord[0])
wy = float(wp_coord[1])
wz = float(wp_coord[2])
if wp_idx == self.current_wp_idx:
# ACTIVE GOAL: Glowing floating target orb + vertical laser beacon column + ground target ring
# 1. Floating target orb (sphere)
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([wx, wy, wz + 0.35], dtype=np.float64)
orb_size = np.array([0.18, 0.18, 0.18], dtype=np.float64)
orb_rgba = np.array([0.98, 0.75, 0.15, 0.90], dtype=np.float32) # Vibrant Gold
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
# 2. Vertical laser beacon column (cylinder)
if scene.ngeom < scene.maxgeom:
dst_beam = scene.geoms[scene.ngeom]
beam_pos = np.array([wx, wy, wz + 0.60], dtype=np.float64)
beam_size = np.array([0.05, 0.60, 0.0], dtype=np.float64)
beam_rgba = np.array([0.22, 0.74, 0.97, 0.55], dtype=np.float32) # Glowing Cyan
mujoco.mjv_initGeom(dst_beam, mujoco.mjtGeom.mjGEOM_CYLINDER, beam_size, beam_pos, identity_mat, beam_rgba)
dst_beam.transparent = 1
scene.ngeom += 1
# 3. Ground target landing pad ring
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
ring_size = np.array([float(wp_radius) * 0.6, 0.02, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.45], dtype=np.float32)
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
elif wp_idx < self.current_wp_idx:
# COMPLETED GOAL: Subtle emerald disc on surface
dst_comp = scene.geoms[scene.ngeom]
comp_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
comp_size = np.array([0.25, 0.015, 0.0], dtype=np.float64)
comp_rgba = np.array([0.10, 0.85, 0.35, 0.40], dtype=np.float32) # Emerald Green
mujoco.mjv_initGeom(dst_comp, mujoco.mjtGeom.mjGEOM_CYLINDER, comp_size, comp_pos, identity_mat, comp_rgba)
dst_comp.transparent = 1
scene.ngeom += 1
else:
# UPCOMING GOAL: Translucent waypoint marker disc
dst_up = scene.geoms[scene.ngeom]
up_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
up_size = np.array([0.16, 0.015, 0.0], dtype=np.float64)
up_rgba = np.array([0.60, 0.65, 0.80, 0.25], dtype=np.float32) # Translucent Silver
mujoco.mjv_initGeom(dst_up, mujoco.mjtGeom.mjGEOM_CYLINDER, up_size, up_pos, identity_mat, up_rgba)
dst_up.transparent = 1
scene.ngeom += 1
# Inject translucent ghost avatars if provided
if ghost_datas:
scene = self._renderer.scene
for g_idx, g_data in enumerate(ghost_datas):
color = GHOST_PALETTE[g_idx % len(GHOST_PALETTE)]
if ghost_alpha is not None:
color = color.copy()
color[3] = float(ghost_alpha)
for gid in self.humanoid_geoms:
if scene.ngeom >= scene.maxgeom:
break
dst = scene.geoms[scene.ngeom]
pos = g_data.geom_xpos[gid].astype(np.float64)
size = self.model.geom_size[gid].astype(np.float64)
mat = g_data.geom_xmat[gid].astype(np.float64)
geom_type = int(self.model.geom_type[gid])
mujoco.mjv_initGeom(dst, geom_type, size, pos, mat, color)
dst.transparent = 1
scene.ngeom += 1
return self._renderer.render()
|