Commit ef90154
committed
ocean/craftax: update_mobs early-out on dead mob slots
The five move_* helpers (melee/passive/ranged mobs + mob/player
projectiles) now return immediately when mask=false. JAX's branchless
"compute-then-mask" pattern is pointless on CPU: dead slots' output
never feeds observations, rewards, or mob_map, so skipping the body
and the RNG draws is semantically equivalent.
Defining CRAFTAX_JAX_PARITY at build time restores the branchless
slow path for bitwise replay against JAX (required by
tests/craftax_parity.py). Default build uses the early-out.
Also drops craftax_step_jax_index(player_level, NUM_LEVELS) clamps at
the top of each move_* -- state->player_level is maintained in
[0, NUM_LEVELS-1] by change_floor_native (explicit bounds checks) and
by the worldgen init. Six redundant clamps per step eliminated.
Measurements (single-thread, random actions, pool=1024):
update_mobs phase: 1.392 us -> 0.285 us (4.88x)
full c_step: 2.35 us -> 1.22 us
1-thread sim SPS: 425K -> 819K (1.93x)
16-thread sim SPS: 5.53M -> 10.04M (1.82x)
training SPS: 506K -> 544K (+7%)
Parity test with CRAFTAX_JAX_PARITY defined passes 8 seeds * 1000
steps over 27 terminals. Without the flag, parity diverges at the
first mob death -- by design.1 parent 93cfb01 commit ef90154
1 file changed
Lines changed: 29 additions & 25 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
600 | 600 | | |
601 | 601 | | |
602 | 602 | | |
603 | | - | |
604 | | - | |
605 | | - | |
606 | | - | |
| 603 | + | |
| 604 | + | |
| 605 | + | |
| 606 | + | |
| 607 | + | |
| 608 | + | |
| 609 | + | |
| 610 | + | |
| 611 | + | |
607 | 612 | | |
608 | 613 | | |
609 | | - | |
610 | 614 | | |
611 | 615 | | |
612 | 616 | | |
| |||
729 | 733 | | |
730 | 734 | | |
731 | 735 | | |
732 | | - | |
733 | | - | |
734 | | - | |
735 | | - | |
| 736 | + | |
| 737 | + | |
| 738 | + | |
| 739 | + | |
| 740 | + | |
736 | 741 | | |
737 | 742 | | |
738 | | - | |
739 | 743 | | |
740 | 744 | | |
741 | 745 | | |
| |||
794 | 798 | | |
795 | 799 | | |
796 | 800 | | |
797 | | - | |
798 | | - | |
799 | | - | |
800 | | - | |
| 801 | + | |
| 802 | + | |
| 803 | + | |
| 804 | + | |
| 805 | + | |
801 | 806 | | |
802 | 807 | | |
803 | | - | |
804 | 808 | | |
805 | 809 | | |
806 | 810 | | |
| |||
932 | 936 | | |
933 | 937 | | |
934 | 938 | | |
935 | | - | |
936 | | - | |
937 | | - | |
938 | | - | |
| 939 | + | |
| 940 | + | |
| 941 | + | |
| 942 | + | |
| 943 | + | |
939 | 944 | | |
940 | 945 | | |
941 | 946 | | |
942 | 947 | | |
943 | 948 | | |
944 | 949 | | |
945 | | - | |
946 | 950 | | |
947 | 951 | | |
948 | 952 | | |
| |||
1009 | 1013 | | |
1010 | 1014 | | |
1011 | 1015 | | |
1012 | | - | |
1013 | | - | |
1014 | | - | |
1015 | | - | |
| 1016 | + | |
| 1017 | + | |
| 1018 | + | |
| 1019 | + | |
| 1020 | + | |
1016 | 1021 | | |
1017 | 1022 | | |
1018 | 1023 | | |
1019 | 1024 | | |
1020 | 1025 | | |
1021 | 1026 | | |
1022 | | - | |
1023 | 1027 | | |
1024 | 1028 | | |
1025 | 1029 | | |
| |||
0 commit comments