flip_landing_reward_wrapper.py
35.5 KB · 1262 lines · python Raw
1
2 from __future__ import annotations
3
4 from collections.abc import Mapping
5 import math
6
7 import gymnasium as gym
8 import numpy as np
9
10
11
12 DEFAULT_REWARD_CONFIG = {
13 # ------------------------------------------------------------------
14 # 1. Task-stage definitions
15 #
16 # These values define when the wrapper considers the flip and
17 # post-flip recovery stages complete.
18 # ------------------------------------------------------------------
19
20 "required_rotations": 1.0,
21 "rotation_direction": 1,
22 "upright_tolerance_radians": 0.30,
23 "recovery_angular_velocity_tolerance": 0.50,
24
25 # ------------------------------------------------------------------
26 # 2. Original LunarLander reward contribution
27 #
28 # The custom reward is added to a weighted version of Gymnasium's
29 # original LunarLander reward.
30 # ------------------------------------------------------------------
31
32 "pre_flip_original_reward_weight": 0.25,
33 "post_flip_original_reward_weight": 1.0,
34
35 # ------------------------------------------------------------------
36 # 3. Positive stage and terminal rewards
37 #
38 # Progress is paid only when the maximum achieved rotation progress
39 # increases. The other rewards are one-off stage rewards.
40 # ------------------------------------------------------------------
41
42 "rotation_progress_bonus": 150.0,
43 "flip_completion_bonus": 200.0,
44 "recovery_bonus": 100.0,
45 "flip_landing_bonus": 750.0,
46
47 # ------------------------------------------------------------------
48 # 4. Terminal failure penalties
49 # ------------------------------------------------------------------
50
51 "landing_without_flip_penalty": 300.0,
52 "no_flip_terminal_penalty": 300.0,
53 "failed_landing_penalty": 150.0,
54 "outside_zone_landing_penalty": 1_000.0,
55 "in_zone_crash_penalty": 1_200.0,
56
57 # ------------------------------------------------------------------
58 # 5. Potential-difference shaping controls
59 #
60 # These control how changes in post-flip state quality are converted
61 # into a dense per-step reward.
62 # ------------------------------------------------------------------
63
64 "post_flip_shaping_weight": 1.0,
65 "post_flip_shaping_gamma": 0.995,
66 "post_flip_shaping_clip": 10.0,
67
68 # ------------------------------------------------------------------
69 # 6. Post-flip state-quality weights
70 #
71 # These determine the relative importance of position, velocity,
72 # attitude, angular velocity and leg contact in the potential.
73 # ------------------------------------------------------------------
74
75 "post_flip_center_weight": 50.0,
76 "post_flip_horizontal_speed_weight": 30.0,
77 "post_flip_vertical_speed_weight": 30.0,
78 "post_flip_angle_weight": 20.0,
79 "post_flip_angular_speed_weight": 10.0,
80 "post_flip_leg_contact_weight": 10.0,
81
82 # ------------------------------------------------------------------
83 # 7. Horizontal landing-zone guidance
84 # ------------------------------------------------------------------
85
86 "landing_zone_half_width": 0.20,
87 "post_flip_zone_excess_weight": 200.0,
88 "post_flip_target_vx_gain": 0.50,
89 "post_flip_max_target_vx": 0.35,
90 "post_flip_horizontal_deadband": 0.08,
91
92 # ------------------------------------------------------------------
93 # 8. Vertical descent and soft-touchdown guidance
94 # ------------------------------------------------------------------
95
96 "post_flip_target_vy_high": -0.45,
97 "post_flip_target_vy_near_ground": -0.12,
98 "near_ground_height": 0.60,
99 "safe_touchdown_vertical_speed": 0.18,
100 "near_ground_overspeed_weight": 120.0,
101 }
102
103
104 class FlipLandingRewardWrapper(gym.Wrapper):
105 """
106 Stage-aware LunarLander objective:
107
108 1. Complete one full rotation in the selected direction.
109 2. Arrest the spin and recover to an upright attitude.
110 3. Re-centre, stabilise and land safely.
111
112 The original eight-value observation is extended with:
113 - signed rotation progress in [-1, 1];
114 - a flip-completed flag;
115 - a post-flip-recovery-completed flag.
116
117 Post-flip recovery uses a potential-difference reward based on
118 horizontal position, horizontal speed, attitude, angular speed,
119 and leg contact. This rewards improvement rather than repeatedly
120 paying the agent simply for occupying a favourable state.
121 """
122
123 def __init__(
124 self,
125 env: gym.Env,
126 *,
127 reward_config: (
128 Mapping[str, float | int]
129 | None
130 ) = None,
131 ):
132 """
133 Initialise the reward wrapper.
134
135 DEFAULT_REWARD_CONFIG is the single source of default values.
136 `reward_config` may override any subset of those values for a
137 particular curriculum phase.
138 """
139
140 super().__init__(env)
141
142 supplied_config = dict(
143 reward_config or {}
144 )
145
146 # Reject misspelled or obsolete parameter names. Without this
147 # check, an accidental typo could silently leave the default
148 # value active.
149 unknown_parameters = (
150 set(supplied_config)
151 - set(DEFAULT_REWARD_CONFIG)
152 )
153
154 if unknown_parameters:
155 raise ValueError(
156 "Unknown reward parameters: "
157 f"{sorted(unknown_parameters)}"
158 )
159
160 # Make a new dictionary so the global defaults are never mutated.
161 config = {
162 **DEFAULT_REWARD_CONFIG,
163 **supplied_config,
164 }
165
166 # ==============================================================
167 # 1. Task-stage definitions
168 # ==============================================================
169
170 self.required_rotations = float(
171 config["required_rotations"]
172 )
173
174 self.rotation_direction = int(
175 config["rotation_direction"]
176 )
177
178 self.upright_tolerance_radians = float(
179 config[
180 "upright_tolerance_radians"
181 ]
182 )
183
184 self.recovery_angular_velocity_tolerance = float(
185 config[
186 "recovery_angular_velocity_tolerance"
187 ]
188 )
189
190 # ==============================================================
191 # 2. Original LunarLander reward contribution
192 # ==============================================================
193
194 self.pre_flip_original_reward_weight = float(
195 config[
196 "pre_flip_original_reward_weight"
197 ]
198 )
199
200 self.post_flip_original_reward_weight = float(
201 config[
202 "post_flip_original_reward_weight"
203 ]
204 )
205
206 # ==============================================================
207 # 3. Positive stage and terminal rewards
208 # ==============================================================
209
210 self.rotation_progress_bonus = float(
211 config[
212 "rotation_progress_bonus"
213 ]
214 )
215
216 self.flip_completion_bonus = float(
217 config[
218 "flip_completion_bonus"
219 ]
220 )
221
222 self.recovery_bonus = float(
223 config["recovery_bonus"]
224 )
225
226 self.flip_landing_bonus = float(
227 config["flip_landing_bonus"]
228 )
229
230 # ==============================================================
231 # 4. Terminal failure penalties
232 # ==============================================================
233
234 self.landing_without_flip_penalty = float(
235 config[
236 "landing_without_flip_penalty"
237 ]
238 )
239
240 self.no_flip_terminal_penalty = float(
241 config[
242 "no_flip_terminal_penalty"
243 ]
244 )
245
246 self.failed_landing_penalty = float(
247 config[
248 "failed_landing_penalty"
249 ]
250 )
251
252 self.outside_zone_landing_penalty = float(
253 config[
254 "outside_zone_landing_penalty"
255 ]
256 )
257
258 self.in_zone_crash_penalty = float(
259 config[
260 "in_zone_crash_penalty"
261 ]
262 )
263
264 # ==============================================================
265 # 5. Potential-difference shaping controls
266 # ==============================================================
267
268 self.post_flip_shaping_weight = float(
269 config[
270 "post_flip_shaping_weight"
271 ]
272 )
273
274 self.post_flip_shaping_gamma = float(
275 config[
276 "post_flip_shaping_gamma"
277 ]
278 )
279
280 self.post_flip_shaping_clip = float(
281 config[
282 "post_flip_shaping_clip"
283 ]
284 )
285
286 # ==============================================================
287 # 6. Post-flip state-quality weights
288 # ==============================================================
289
290 self.post_flip_center_weight = float(
291 config[
292 "post_flip_center_weight"
293 ]
294 )
295
296 self.post_flip_horizontal_speed_weight = float(
297 config[
298 "post_flip_horizontal_speed_weight"
299 ]
300 )
301
302 self.post_flip_vertical_speed_weight = float(
303 config[
304 "post_flip_vertical_speed_weight"
305 ]
306 )
307
308 self.post_flip_angle_weight = float(
309 config[
310 "post_flip_angle_weight"
311 ]
312 )
313
314 self.post_flip_angular_speed_weight = float(
315 config[
316 "post_flip_angular_speed_weight"
317 ]
318 )
319
320 self.post_flip_leg_contact_weight = float(
321 config[
322 "post_flip_leg_contact_weight"
323 ]
324 )
325
326 # ==============================================================
327 # 7. Horizontal landing-zone guidance
328 # ==============================================================
329
330 self.landing_zone_half_width = float(
331 config[
332 "landing_zone_half_width"
333 ]
334 )
335
336 self.post_flip_zone_excess_weight = float(
337 config[
338 "post_flip_zone_excess_weight"
339 ]
340 )
341
342 self.post_flip_target_vx_gain = float(
343 config[
344 "post_flip_target_vx_gain"
345 ]
346 )
347
348 self.post_flip_max_target_vx = float(
349 config[
350 "post_flip_max_target_vx"
351 ]
352 )
353
354 self.post_flip_horizontal_deadband = float(
355 config[
356 "post_flip_horizontal_deadband"
357 ]
358 )
359
360 # ==============================================================
361 # 8. Vertical descent and soft-touchdown guidance
362 # ==============================================================
363
364 self.post_flip_target_vy_high = float(
365 config[
366 "post_flip_target_vy_high"
367 ]
368 )
369
370 self.post_flip_target_vy_near_ground = float(
371 config[
372 "post_flip_target_vy_near_ground"
373 ]
374 )
375
376 self.near_ground_height = float(
377 config[
378 "near_ground_height"
379 ]
380 )
381
382 self.safe_touchdown_vertical_speed = float(
383 config[
384 "safe_touchdown_vertical_speed"
385 ]
386 )
387
388 self.near_ground_overspeed_weight = float(
389 config[
390 "near_ground_overspeed_weight"
391 ]
392 )
393
394 # ==============================================================
395 # Configuration validation
396 # ==============================================================
397
398 if self.required_rotations <= 0.0:
399 raise ValueError(
400 "required_rotations must be positive."
401 )
402
403 if self.rotation_direction not in (
404 -1,
405 1,
406 ):
407 raise ValueError(
408 "rotation_direction must be "
409 "either -1 or 1."
410 )
411
412 if not (
413 0.0
414 < self.upright_tolerance_radians
415 <= math.pi
416 ):
417 raise ValueError(
418 "upright_tolerance_radians must "
419 "be in the interval (0, pi]."
420 )
421
422 if (
423 self.recovery_angular_velocity_tolerance
424 < 0.0
425 ):
426 raise ValueError(
427 "recovery_angular_velocity_tolerance "
428 "must not be negative."
429 )
430
431 if not (
432 0.0
433 <= self.post_flip_shaping_gamma
434 <= 1.0
435 ):
436 raise ValueError(
437 "post_flip_shaping_gamma must "
438 "be in [0, 1]."
439 )
440
441 if self.post_flip_shaping_clip <= 0.0:
442 raise ValueError(
443 "post_flip_shaping_clip must "
444 "be positive."
445 )
446
447 if not (
448 0.0
449 < self.landing_zone_half_width
450 < 1.0
451 ):
452 raise ValueError(
453 "landing_zone_half_width must "
454 "be between 0 and 1."
455 )
456
457 if self.post_flip_max_target_vx <= 0.0:
458 raise ValueError(
459 "post_flip_max_target_vx must "
460 "be positive."
461 )
462
463 if self.near_ground_height <= 0.0:
464 raise ValueError(
465 "near_ground_height must be "
466 "positive."
467 )
468
469 # All parameters in this collection represent rewards,
470 # penalties, weights, tolerances or magnitudes. Negative values
471 # would invert their intended meaning.
472 non_negative_parameters = {
473 "pre_flip_original_reward_weight": (
474 self.pre_flip_original_reward_weight
475 ),
476 "post_flip_original_reward_weight": (
477 self.post_flip_original_reward_weight
478 ),
479 "rotation_progress_bonus": (
480 self.rotation_progress_bonus
481 ),
482 "flip_completion_bonus": (
483 self.flip_completion_bonus
484 ),
485 "recovery_bonus": (
486 self.recovery_bonus
487 ),
488 "flip_landing_bonus": (
489 self.flip_landing_bonus
490 ),
491 "landing_without_flip_penalty": (
492 self.landing_without_flip_penalty
493 ),
494 "no_flip_terminal_penalty": (
495 self.no_flip_terminal_penalty
496 ),
497 "failed_landing_penalty": (
498 self.failed_landing_penalty
499 ),
500 "outside_zone_landing_penalty": (
501 self.outside_zone_landing_penalty
502 ),
503 "in_zone_crash_penalty": (
504 self.in_zone_crash_penalty
505 ),
506 "post_flip_shaping_weight": (
507 self.post_flip_shaping_weight
508 ),
509 "post_flip_center_weight": (
510 self.post_flip_center_weight
511 ),
512 "post_flip_horizontal_speed_weight": (
513 self.post_flip_horizontal_speed_weight
514 ),
515 "post_flip_vertical_speed_weight": (
516 self.post_flip_vertical_speed_weight
517 ),
518 "post_flip_angle_weight": (
519 self.post_flip_angle_weight
520 ),
521 "post_flip_angular_speed_weight": (
522 self.post_flip_angular_speed_weight
523 ),
524 "post_flip_leg_contact_weight": (
525 self.post_flip_leg_contact_weight
526 ),
527 "post_flip_zone_excess_weight": (
528 self.post_flip_zone_excess_weight
529 ),
530 "post_flip_target_vx_gain": (
531 self.post_flip_target_vx_gain
532 ),
533 "post_flip_horizontal_deadband": (
534 self.post_flip_horizontal_deadband
535 ),
536 "safe_touchdown_vertical_speed": (
537 self.safe_touchdown_vertical_speed
538 ),
539 "near_ground_overspeed_weight": (
540 self.near_ground_overspeed_weight
541 ),
542 }
543
544 for (
545 parameter_name,
546 parameter_value,
547 ) in non_negative_parameters.items():
548 if parameter_value < 0.0:
549 raise ValueError(
550 f"{parameter_name} must not "
551 "be negative."
552 )
553
554 # Downward velocity is represented by a negative value.
555 if self.post_flip_target_vy_high > 0.0:
556 raise ValueError(
557 "post_flip_target_vy_high must "
558 "be zero or negative."
559 )
560
561 if (
562 self.post_flip_target_vy_near_ground
563 > 0.0
564 ):
565 raise ValueError(
566 "post_flip_target_vy_near_ground "
567 "must be zero or negative."
568 )
569
570 # For example, -0.12 is slower than -0.45.
571 if (
572 self.post_flip_target_vy_near_ground
573 < self.post_flip_target_vy_high
574 ):
575 raise ValueError(
576 "The near-ground target vertical "
577 "speed must be slower than the "
578 "high-altitude target."
579 )
580
581 if (
582 self.post_flip_horizontal_deadband
583 > self.landing_zone_half_width
584 ):
585 raise ValueError(
586 "post_flip_horizontal_deadband "
587 "must not exceed "
588 "landing_zone_half_width."
589 )
590
591 # Convert the configured number of rotations to radians.
592 self.required_rotation = (
593 2.0
594 * math.pi
595 * self.required_rotations
596 )
597
598 # ==============================================================
599 # Per-episode state
600 #
601 # These values are reset at the start of every episode. They are
602 # not reward hyperparameters.
603 # ==============================================================
604
605 self.previous_angle = 0.0
606 self.cumulative_rotation = 0.0
607 self.maximum_rotation_progress = 0.0
608 self.flip_completed = False
609 self.recovery_completed = False
610
611 self.previous_post_flip_potential: (
612 float | None
613 ) = None
614
615 # ==============================================================
616 # Extended observation space
617 #
618 # Append:
619 # 1. signed rotation progress;
620 # 2. flip-completed flag;
621 # 3. recovery-completed flag.
622 # ==============================================================
623
624 base_space = self.env.observation_space
625
626 if not isinstance(
627 base_space,
628 gym.spaces.Box,
629 ):
630 raise TypeError(
631 "The wrapped environment must "
632 "have a Box observation space."
633 )
634
635 extra_low = np.asarray(
636 [
637 -1.0,
638 0.0,
639 0.0,
640 ],
641 dtype=np.float32,
642 )
643
644 extra_high = np.asarray(
645 [
646 1.0,
647 1.0,
648 1.0,
649 ],
650 dtype=np.float32,
651 )
652
653 self.observation_space = (
654 gym.spaces.Box(
655 low=np.concatenate(
656 [
657 base_space.low.astype(
658 np.float32
659 ),
660 extra_low,
661 ]
662 ),
663 high=np.concatenate(
664 [
665 base_space.high.astype(
666 np.float32
667 ),
668 extra_high,
669 ]
670 ),
671 dtype=np.float32,
672 )
673 )
674
675 @staticmethod
676 def _wrap_angle(angle: float) -> float:
677 """Wrap an angle to [-pi, pi]."""
678
679 return float(
680 np.arctan2(
681 np.sin(angle),
682 np.cos(angle),
683 )
684 )
685
686 def _signed_rotation_progress(self) -> float:
687 return float(
688 np.clip(
689 self.cumulative_rotation
690 / self.required_rotation,
691 -1.0,
692 1.0,
693 )
694 )
695
696 def _augment_observation(
697 self,
698 observation: np.ndarray,
699 ) -> np.ndarray:
700 return np.concatenate(
701 [
702 np.asarray(
703 observation,
704 dtype=np.float32,
705 ),
706 np.asarray(
707 [
708 self._signed_rotation_progress(),
709 float(self.flip_completed),
710 float(self.recovery_completed),
711 ],
712 dtype=np.float32,
713 ),
714 ]
715 )
716
717 def _post_flip_potential(
718 self,
719 observation: np.ndarray,
720 ) -> float:
721 """
722 Higher values represent safer post-flip states.
723
724 LunarLander-v3 observation positions:
725 0 x position, 2 x velocity, 4 angle, 5 angular velocity,
726 6 left-leg contact and 7 right-leg contact.
727 """
728
729 x_position = float(
730 observation[0]
731 )
732
733 y_position = float(
734 observation[1]
735 )
736
737 horizontal_speed = float(
738 observation[2]
739 )
740
741 vertical_speed = float(
742 observation[3]
743 )
744
745 height_ratio = float(
746 np.clip(
747 y_position
748 / self.near_ground_height,
749 0.0,
750 1.0,
751 )
752 )
753
754 target_vertical_speed = (
755 self.post_flip_target_vy_near_ground
756 + height_ratio
757 * (
758 self.post_flip_target_vy_high
759 - self.post_flip_target_vy_near_ground
760 )
761 )
762
763 vertical_speed_error = abs(
764 vertical_speed
765 - target_vertical_speed
766 )
767
768 wrapped_angle = abs(
769 self._wrap_angle(
770 float(observation[4])
771 )
772 )
773
774 angular_speed = abs(
775 float(observation[5])
776 )
777
778 leg_contacts = (
779 float(observation[6] > 0.5)
780 + float(observation[7] > 0.5)
781 )
782
783 # Distance beyond the two landing-zone flags.
784 zone_excess = max(
785 abs(x_position)
786 - self.landing_zone_half_width,
787 0.0,
788 )
789
790 # When left of the pad, target positive horizontal
791 # velocity. When right of the pad, target negative
792 # horizontal velocity. Close to the centre, target
793 # velocity approaches zero.
794 horizontal_guidance_factor = float(
795 np.clip(
796 y_position / 0.50,
797 0.0,
798 1.0,
799 )
800 )
801
802 horizontal_error = max(
803 abs(x_position)
804 - self.post_flip_horizontal_deadband,
805 0.0,
806 )
807
808 target_horizontal_speed = float(
809 np.clip(
810 -self.post_flip_target_vx_gain
811 * x_position
812 * horizontal_guidance_factor,
813 -self.post_flip_max_target_vx,
814 self.post_flip_max_target_vx,
815 )
816 )
817
818 horizontal_speed_error = abs(
819 horizontal_speed
820 - target_horizontal_speed
821 )
822
823 # Increase the importance of entering the landing
824 # zone as the lander approaches the ground.
825 near_ground_factor = float(
826 np.clip(
827 1.0
828 - max(y_position, 0.0)
829 / 0.75,
830 0.0,
831 1.0,
832 )
833 )
834
835 effective_zone_weight = (
836 self.post_flip_zone_excess_weight
837 * (1.0 + near_ground_factor)
838 )
839
840 return float(
841 -self.post_flip_center_weight
842 * horizontal_error
843
844 -effective_zone_weight
845 * zone_excess
846
847 -self.post_flip_horizontal_speed_weight
848 * horizontal_speed_error
849
850 -self.post_flip_vertical_speed_weight
851 * vertical_speed_error
852
853 -self.post_flip_angle_weight
854 * wrapped_angle
855
856 -self.post_flip_angular_speed_weight
857 * angular_speed
858
859 +self.post_flip_leg_contact_weight
860 * leg_contacts
861 )
862
863 def reset(
864 self,
865 *,
866 seed: int | None = None,
867 options: dict | None = None,
868 ):
869 observation, info = self.env.reset(
870 seed=seed,
871 options=options,
872 )
873
874 self.previous_angle = float(observation[4])
875 self.cumulative_rotation = 0.0
876 self.maximum_rotation_progress = 0.0
877 self.flip_completed = False
878 self.recovery_completed = False
879 self.previous_post_flip_potential = None
880
881 info = dict(info)
882 info.update(
883 {
884 "flip_completed": False,
885 "recovery_completed": False,
886 "rotation_progress": 0.0,
887 "rotations_completed": 0.0,
888 "rotation_progress_reward": 0.0,
889 "flip_completion_reward": 0.0,
890 "recovery_reward": 0.0,
891 "post_flip_shaping_reward": 0.0,
892 "flip_landing_bonus": 0.0,
893 "terminal_adjustment": 0.0,
894 }
895 )
896
897 return self._augment_observation(observation), info
898
899 def step(self, action):
900 # 1. Advance the original LunarLander environment.
901 (
902 observation,
903 original_reward,
904 terminated,
905 truncated,
906 info,
907 ) = self.env.step(action)
908
909 # 2. Accumulate directed angular movement and reward new flip progress.
910 current_angle = float(observation[4])
911 angular_velocity = float(observation[5])
912
913 angle_change = self._wrap_angle(
914 current_angle - self.previous_angle
915 )
916 self.previous_angle = current_angle
917
918 directed_change = (
919 self.rotation_direction * angle_change
920 )
921 self.cumulative_rotation += directed_change
922
923 rotation_progress = float(
924 np.clip(
925 self.cumulative_rotation
926 / self.required_rotation,
927 0.0,
928 1.0,
929 )
930 )
931
932 new_progress = max(
933 0.0,
934 rotation_progress
935 - self.maximum_rotation_progress,
936 )
937 progress_reward = (
938 self.rotation_progress_bonus
939 * new_progress
940 )
941 self.maximum_rotation_progress = max(
942 self.maximum_rotation_progress,
943 rotation_progress,
944 )
945
946 completion_reward = 0.0
947
948 # 3. Detect the first completed full rotation.
949 if (
950 rotation_progress >= 1.0
951 and not self.flip_completed
952 ):
953 self.flip_completed = True
954 completion_reward = (
955 self.flip_completion_bonus
956 )
957
958 # 4. Detect upright post-flip recovery.
959 wrapped_angle = abs(
960 self._wrap_angle(current_angle)
961 )
962 upright = (
963 wrapped_angle
964 <= self.upright_tolerance_radians
965 )
966
967 recovery_reward = 0.0
968
969 if (
970 self.flip_completed
971 and not self.recovery_completed
972 and upright
973 and abs(angular_velocity)
974 <= self.recovery_angular_velocity_tolerance
975 ):
976 self.recovery_completed = True
977 recovery_reward = self.recovery_bonus
978
979
980 post_flip_shaping_reward = 0.0
981 post_flip_potential = None
982 near_ground_overspeed_penalty = 0.0
983
984 # 5. Apply an explicit near-ground descen-overspeed penalty.
985 if self.flip_completed:
986 y_position = float(
987 observation[1]
988 )
989
990 vertical_speed = float(
991 observation[3]
992 )
993
994 near_ground_factor = float(
995 np.clip(
996 1.0
997 - max(y_position, 0.0)
998 / self.near_ground_height,
999 0.0,
1000 1.0,
1001 )
1002 )
1003
1004 downward_speed = max(
1005 -vertical_speed,
1006 0.0,
1007 )
1008
1009 descent_overspeed = max(
1010 downward_speed
1011 - self.safe_touchdown_vertical_speed,
1012 0.0,
1013 )
1014
1015 near_ground_overspeed_penalty = (
1016 self.near_ground_overspeed_weight
1017 * near_ground_factor
1018 * descent_overspeed**2
1019 )
1020
1021 # 6. Apply continuous post-flip landing-quality shaping.
1022 if self.flip_completed:
1023 post_flip_potential = (
1024 self._post_flip_potential(observation)
1025 )
1026
1027 if self.previous_post_flip_potential is None:
1028 self.previous_post_flip_potential = (
1029 post_flip_potential
1030 )
1031 else:
1032 potential_difference = (
1033 self.post_flip_shaping_gamma
1034 * post_flip_potential
1035 - self.previous_post_flip_potential
1036 )
1037 post_flip_shaping_reward = float(
1038 np.clip(
1039 self.post_flip_shaping_weight
1040 * potential_difference,
1041 -self.post_flip_shaping_clip,
1042 self.post_flip_shaping_clip,
1043 )
1044 )
1045 self.previous_post_flip_potential = (
1046 post_flip_potential
1047 )
1048
1049 left_leg_contact = bool(observation[6] > 0.5)
1050 right_leg_contact = bool(observation[7] > 0.5)
1051
1052 x_position = float(
1053 observation[0]
1054 )
1055
1056 inside_landing_zone = bool(
1057 abs(x_position)
1058 <= self.landing_zone_half_width
1059 )
1060
1061 stable_landing = bool(
1062 terminated
1063 and float(original_reward) > 0.0
1064 and left_leg_contact
1065 and right_leg_contact
1066 and upright
1067 )
1068
1069 # "Safe landing" now means stable and
1070 # between the landing-zone boundaries.
1071 landed_safely = bool(
1072 stable_landing
1073 and inside_landing_zone
1074 )
1075
1076 episode_finished = bool(
1077 terminated or truncated
1078 )
1079
1080 # 7. Classify the terminal outcome and apply its one-off adjustment.
1081 terminal_adjustment = 0.0
1082 outcome = "in_progress"
1083
1084 # Complete success:
1085 # flipped, settled safely and landed inside the zone.
1086 if (
1087 self.flip_completed
1088 and landed_safely
1089 ):
1090 terminal_adjustment = (
1091 self.flip_landing_bonus
1092 )
1093
1094 outcome = (
1095 "flip_and_safe_landing"
1096 )
1097
1098 elif episode_finished:
1099 # The episode ended before completing the flip.
1100 if not self.flip_completed:
1101 if stable_landing:
1102 terminal_adjustment = -(
1103 self.landing_without_flip_penalty
1104 )
1105
1106 outcome = (
1107 "stable_landing_without_flip"
1108 )
1109
1110 else:
1111 terminal_adjustment = -(
1112 self.no_flip_terminal_penalty
1113 )
1114
1115 outcome = (
1116 "episode_ended_without_flip"
1117 )
1118
1119 # Step E:
1120 # The flip was completed and the lander reached
1121 # the zone, but it crashed instead of settling.
1122 elif (
1123 terminated
1124 and inside_landing_zone
1125 and not stable_landing
1126 ):
1127 terminal_adjustment = -(
1128 self.in_zone_crash_penalty
1129 )
1130
1131 outcome = (
1132 "flip_and_crash_in_zone"
1133 )
1134
1135 # A stable landing occurred, but outside the flags.
1136 elif (
1137 stable_landing
1138 and not inside_landing_zone
1139 ):
1140 terminal_adjustment = -(
1141 self.outside_zone_landing_penalty
1142 )
1143
1144 outcome = (
1145 "flip_and_stable_landing_"
1146 "outside_zone"
1147 )
1148
1149 # Other post-flip failures, such as crashing
1150 # outside the zone or timing out.
1151 else:
1152 terminal_adjustment = -(
1153 self.failed_landing_penalty
1154 )
1155
1156 outcome = (
1157 "flip_but_failed_landing"
1158 )
1159
1160
1161 original_weight = (
1162 self.post_flip_original_reward_weight
1163 if self.flip_completed
1164 else self.pre_flip_original_reward_weight
1165 )
1166
1167 # 8. Combine the original reward, dense shaping and terminal outcome.
1168 modified_reward = (
1169 original_weight * float(original_reward)
1170 + progress_reward
1171 + completion_reward
1172 + recovery_reward
1173 + post_flip_shaping_reward
1174 - near_ground_overspeed_penalty
1175 + terminal_adjustment
1176 )
1177
1178 info = dict(info)
1179 info.update(
1180 {
1181 "original_reward": float(original_reward),
1182 "rotation_progress": rotation_progress,
1183 "rotation_progress_reward": progress_reward,
1184 "flip_completion_reward": completion_reward,
1185 "recovery_reward": recovery_reward,
1186 "post_flip_shaping_reward": (
1187 post_flip_shaping_reward
1188 ),
1189 "post_flip_potential": post_flip_potential,
1190 "cumulative_rotation": (
1191 self.cumulative_rotation
1192 ),
1193 "rotations_completed": (
1194 self.cumulative_rotation
1195 / (2.0 * math.pi)
1196 ),
1197 "flip_completed": self.flip_completed,
1198 "recovery_completed": (
1199 self.recovery_completed
1200 ),
1201 "landed_safely": landed_safely,
1202 "original_reward_weight": original_weight,
1203 "terminal_adjustment": (
1204 terminal_adjustment
1205 ),
1206 "flip_landing_bonus": (
1207 self.flip_landing_bonus
1208 if outcome
1209 == "flip_and_safe_landing"
1210 else 0.0
1211 ),
1212 "custom_outcome": outcome,
1213 "stable_landing": stable_landing,
1214 "inside_landing_zone": (
1215 inside_landing_zone
1216 ),
1217 "horizontal_position": (
1218 x_position
1219 ),
1220 "near_ground_overspeed_penalty": (
1221 near_ground_overspeed_penalty
1222 ),
1223 "vertical_speed": float(
1224 observation[3]
1225 ),
1226 }
1227 )
1228
1229 return (
1230 self._augment_observation(observation),
1231 modified_reward,
1232 terminated,
1233 truncated,
1234 info,
1235 )
1236
1237
1238 def make_flip_lander(
1239 *,
1240 render_mode: str | None = None,
1241 reward_config: (
1242 Mapping[str, float | int]
1243 | None
1244 ) = None,
1245 ) -> gym.Env:
1246 """
1247 Build the exact environment used for training and evaluation.
1248
1249 Pass the same reward_config to the uploader so the Hub model card
1250 records the environment accurately.
1251 """
1252
1253 base_env = gym.make(
1254 "LunarLander-v3",
1255 render_mode=render_mode,
1256 )
1257
1258 return FlipLandingRewardWrapper(
1259 base_env,
1260 reward_config=reward_config,
1261 )
1262