diff --git a/example/cartpole/discretizer.py b/example/cartpole/discretizer.py index ffe43d6..db9f7d7 100644 --- a/example/cartpole/discretizer.py +++ b/example/cartpole/discretizer.py @@ -33,6 +33,14 @@ class Action(Enum): class CartpoleDiscretizer(Discretizer): + PREDICATE_TYPES = (Position, Velocity, Angle) + + PREDICATE_ORDER = { + Position: 0, + Velocity: 1, + Angle: 2, + } + def __init__(self): super(CartpoleDiscretizer, self).__init__() @@ -53,7 +61,7 @@ def discretize(self, state: np.ndarray) -> Tuple[Predicate, Predicate, Predicate stuck_velocity_thr = 0.1 standing_angle_thr = 0.0005 - pole_predicate = "" + if -standing_angle_thr < angle < standing_angle_thr: pole_predicate = Angle.STANDING elif angle < 0 and -stuck_velocity_thr < ang_velocity < stuck_velocity_thr: @@ -68,6 +76,11 @@ def discretize(self, state: np.ndarray) -> Tuple[Predicate, Predicate, Predicate pole_predicate = Angle.FALLING_RIGHT elif angle > 0 and ang_velocity < 0: pole_predicate = Angle.STABILIZING_LEFT + else: + raise ValueError( + f"Could not discretize pole state: angle={angle}, " + f"ang_velocity={ang_velocity}" + ) return ( Predicate(pos_predicate), @@ -76,41 +89,146 @@ def discretize(self, state: np.ndarray) -> Tuple[Predicate, Predicate, Predicate ) def state_to_str(self, state: Tuple[Predicate, Predicate, Predicate]) -> str: - return "&".join(str(pred) for pred in state) + """ + Serialize states using a deterministic and unambiguous predicate order. + + The default Predicate.__str__ representation may produce strings such as + RIGHT(), which are ambiguous because both Position and Velocity contain + RIGHT. Therefore, this method explicitly stores the enum type name. + """ + ordered_state = sorted( + state, + key=lambda predicate: self._predicate_sort_key(predicate), + ) + + return "&".join( + self._predicate_to_str(predicate) + for predicate in ordered_state + ) + + def str_to_state(self, state: str) -> Tuple[Predicate, Predicate, Predicate]: + """ + Deserialize a state string independently of predicate order. + + This method does not assume that the string is saved as: + Position & Velocity & Angle + + Instead, each serialized predicate is identified by its enum type. + """ + parsed_predicates = {} - def str_to_state(self, state: str): - pos, vel, angle = state.split("&") - pos_predicate = Position[pos[:-1].split("(")[1]] - mov_predicate = Velocity[vel[:-1].split("(")[1]] - pole_predicate = Angle[angle[:-1].split("(")[1]] + for raw_predicate in state.split("&"): + predicate = self._str_to_predicate(raw_predicate) + + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + if enum_type in parsed_predicates: + raise ValueError( + f"Duplicated predicate type {enum_type.__name__} " + f"in serialized state: {state}" + ) + + parsed_predicates[enum_type] = predicate + + missing_types = [ + enum_type.__name__ + for enum_type in self.PREDICATE_TYPES + if enum_type not in parsed_predicates + ] + + if missing_types: + raise ValueError( + f"Serialized state is missing predicates of type " + f"{missing_types}: {state}" + ) return ( - Predicate(pos_predicate), - Predicate(mov_predicate), - Predicate(pole_predicate), + parsed_predicates[Position], + parsed_predicates[Velocity], + parsed_predicates[Angle], ) + def _predicate_to_str(self, predicate: Predicate) -> str: + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + return f"Predicate({enum_type.__name__}.{enum_value.name})" + + def _predicate_sort_key(self, predicate: Predicate) -> int: + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + if enum_type not in self.PREDICATE_ORDER: + raise ValueError(f"Unknown predicate type: {enum_type}") + + return self.PREDICATE_ORDER[enum_type] + + def _predicate_enum_value(self, predicate: Predicate) -> Enum: + """ + Extract the enum value stored inside a Predicate. + + In the current pgeon Predicate implementation, the enum value is stored + in predicate.name, not in predicate.value. + """ + value = predicate.name + + if not isinstance(value, Enum): + raise ValueError(f"Predicate does not contain an enum name: {predicate}") + + return value + + def _str_to_predicate(self, raw_predicate: str) -> Predicate: + """ + Parse one serialized predicate. + + The parser identifies predicates by enum type and enum member name, so it + does not depend on the order in which predicates appear in the state + string. + """ + raw_predicate = raw_predicate.strip() + + for enum_type in self.PREDICATE_TYPES: + for enum_value in enum_type: + enum_type_name = enum_type.__name__ + enum_member_name = enum_value.name + + valid_patterns = [ + f"Predicate({enum_type_name}.{enum_member_name})", + f"{enum_type_name}.{enum_member_name}", + f"{enum_type_name}({enum_member_name})", + f"{enum_type_name}: {enum_member_name}", + ] + + if raw_predicate in valid_patterns: + return Predicate(enum_value) + + raise ValueError(f"Could not parse predicate from string: {raw_predicate}") + def nearest_state(self, state): og_position, og_velocity, og_angle = state for e in Position: - if [e] != og_position.value: + if Predicate(e) != og_position: yield Predicate(e), og_velocity, og_angle + for e in Velocity: - if [e] != og_velocity.value: + if Predicate(e) != og_velocity: yield og_position, Predicate(e), og_angle + for e in Angle: - if [e] != og_angle.value: + if Predicate(e) != og_angle: yield og_position, og_velocity, Predicate(e) for e in Position: for f in Velocity: for g in Angle: amount_of_equals_to_og = ( - int([e] == og_position.value) - + int([f] == og_velocity.value) - + int([g] == og_angle.value) + int(Predicate(e) == og_position) + + int(Predicate(f) == og_velocity) + + int(Predicate(g) == og_angle) ) + if amount_of_equals_to_og <= 1: yield ( Predicate(e), @@ -123,8 +241,10 @@ def all_actions(self): def get_predicate_space(self): all_tuples = [] + for p in Position: for v in Velocity: for a in Angle: all_tuples.append((p, v, a)) - return all_tuples + + return all_tuples \ No newline at end of file diff --git a/src/pgeon/discretizer.py b/src/pgeon/discretizer.py index 8c5e295..0ca4159 100644 --- a/src/pgeon/discretizer.py +++ b/src/pgeon/discretizer.py @@ -111,20 +111,141 @@ class StateMetadata(BaseModel): class Discretizer(abc.ABC): + PREDICATE_TYPES: Sequence[type[Enum]] = () + @abc.abstractmethod - def discretize(self, non_discrete_state) -> State: ... + def discretize(self, non_discrete_state) -> State: + ... @abc.abstractmethod - def state_to_str(self, state: State) -> str: ... + def state_to_str(self, state: State) -> str: + predicates = self._extract_predicates(state) + + ordered_predicates = sorted( + predicates, + key=lambda predicate: self._predicate_sort_key(predicate), + ) + + return "&".join( + self._predicate_to_str(predicate) + for predicate in ordered_predicates + ) @abc.abstractmethod - def str_to_state(self, state_str: str) -> Predicate: ... + def str_to_state(self, state_str: str) -> State: + if not self.PREDICATE_TYPES: + raise NotImplementedError( + f"{self.__class__.__name__} must define PREDICATE_TYPES " + f"or implement str_to_state manually." + ) + + parsed_predicates = {} + + for raw_predicate in state_str.split("&"): + if not raw_predicate.strip(): + continue + + predicate = self._str_to_predicate(raw_predicate) + + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + if enum_type in parsed_predicates: + raise ValueError( + f"Duplicated predicate type {enum_type.__name__} " + f"in serialized state: {state_str}" + ) + + parsed_predicates[enum_type] = predicate + + missing_types = [ + enum_type.__name__ + for enum_type in self.PREDICATE_TYPES + if enum_type not in parsed_predicates + ] + + if missing_types: + raise ValueError( + f"Serialized state is missing predicates of type " + f"{missing_types}: {state_str}" + ) + + return tuple( + parsed_predicates[enum_type] + for enum_type in self.PREDICATE_TYPES + ) + + def _extract_predicates(self, state: State) -> tuple[Predicate, ...]: + if isinstance(state, PredicateBasedState): + return tuple(state.predicates) + + if isinstance(state, tuple) and all( + isinstance(predicate, Predicate) + for predicate in state + ): + return state + + raise ValueError( + f"Cannot serialize state of type {type(state).__name__}. " + f"Expected PredicateBasedState or tuple[Predicate, ...]." + ) + + def _predicate_to_str(self, predicate: Predicate) -> str: + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + return f"Predicate({enum_type.__name__}.{enum_value.name})" + + def _str_to_predicate(self, raw_predicate: str) -> Predicate: + raw_predicate = raw_predicate.strip() + + if not raw_predicate.startswith("Predicate(") or not raw_predicate.endswith(")"): + raise ValueError(f"Invalid predicate format: {raw_predicate}") + + content = raw_predicate[len("Predicate("):-1] + + if "." not in content: + raise ValueError(f"Invalid predicate content: {content}") + + enum_type_name, enum_member_name = content.split(".", 1) + + for enum_type in self.PREDICATE_TYPES: + if enum_type.__name__ == enum_type_name: + try: + return Predicate(enum_type[enum_member_name]) + except KeyError as exc: + raise ValueError( + f"Unknown enum member {enum_member_name} " + f"for enum type {enum_type_name}" + ) from exc + + raise ValueError(f"Unknown enum type: {enum_type_name}") + + def _predicate_sort_key(self, predicate: Predicate) -> int: + enum_value = self._predicate_enum_value(predicate) + enum_type = type(enum_value) + + if enum_type not in self.PREDICATE_TYPES: + raise ValueError(f"Unknown predicate type: {enum_type}") + + return list(self.PREDICATE_TYPES).index(enum_type) + + def _predicate_enum_value(self, predicate: Predicate) -> Enum: + value = predicate.name + + if not isinstance(value, Enum): + raise ValueError(f"Predicate does not contain an enum name: {predicate}") + + return value @abc.abstractmethod - def nearest_state(self, state) -> Iterator[State]: ... + def nearest_state(self, state) -> Iterator[State]: + ... @abc.abstractmethod - def all_actions(self) -> Sequence[Action]: ... + def all_actions(self) -> Sequence[Action]: + ... @abc.abstractmethod - def get_predicate_space(self) -> Sequence[Predicate]: ... + def get_predicate_space(self) -> Sequence[Predicate]: + ... \ No newline at end of file diff --git a/test/examples/test_cartpole_discretizer.py b/test/examples/test_cartpole_discretizer.py new file mode 100644 index 0000000..63bd883 --- /dev/null +++ b/test/examples/test_cartpole_discretizer.py @@ -0,0 +1,174 @@ +import unittest +from enum import Enum, auto + +from pgeon import Predicate + +from example.cartpole.discretizer import ( + CartpoleDiscretizer, + Position, + Velocity, + Angle, +) + + +class TestCartpoleDiscretizer(unittest.TestCase): + + # ------------------------- + # HAPPY PATHS + # ------------------------- + + def test_str_to_state_handles_predicates_in_expected_order(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + state = discretizer.str_to_state(state_str) + + self.assertEqual( + state, + ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(Angle.STANDING), + ), + ) + + def test_str_to_state_handles_unordered_predicates(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Angle.STANDING)&" + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)" + ) + + state = discretizer.str_to_state(state_str) + + self.assertEqual( + state, + ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(Angle.STANDING), + ), + ) + + def test_state_to_str_uses_consistent_order(self): + discretizer = CartpoleDiscretizer() + + state = ( + Predicate(Angle.STANDING), + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + ) + + state_str = discretizer.state_to_str(state) + + self.assertLess(state_str.index("Position"), state_str.index("Velocity")) + self.assertLess(state_str.index("Velocity"), state_str.index("Angle")) + + def test_state_serialization_round_trip_with_unordered_state(self): + discretizer = CartpoleDiscretizer() + + original_state = ( + Predicate(Angle.FALLING_LEFT), + Predicate(Velocity.LEFT), + Predicate(Position.RIGHT), + ) + + serialized = discretizer.state_to_str(original_state) + loaded_state = discretizer.str_to_state(serialized) + + self.assertEqual( + loaded_state, + ( + Predicate(Position.RIGHT), + Predicate(Velocity.LEFT), + Predicate(Angle.FALLING_LEFT), + ), + ) + + # ------------------------- + # UNHAPPY PATHS + # ------------------------- + + def test_str_to_state_raises_error_when_position_is_missing(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_velocity_is_missing(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_angle_is_missing(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)" + ) + + with self.assertRaises(ValueError): + discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_predicate_type_is_duplicated(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Position.LEFT)&" + "Predicate(Position.RIGHT)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_for_unknown_predicate_type(self): + discretizer = CartpoleDiscretizer() + + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Unknown.VALUE)" + ) + + with self.assertRaises(ValueError): + discretizer.str_to_state(state_str) + + def test_state_to_str_raises_error_for_unknown_predicate_type(self): + discretizer = CartpoleDiscretizer() + + class UnknownPredicateType(Enum): + UNKNOWN = auto() + + state = ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(UnknownPredicateType.UNKNOWN), + ) + + with self.assertRaises(ValueError): + discretizer.state_to_str(state) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/test/pgeon/test_discretizer.py b/test/pgeon/test_discretizer.py index e3f7ab4..57583ef 100644 --- a/test/pgeon/test_discretizer.py +++ b/test/pgeon/test_discretizer.py @@ -1,9 +1,12 @@ import unittest from enum import Enum, auto +from typing import Iterator, Sequence from pydantic import ValidationError from pgeon.discretizer import ( + Action, + Discretizer, Predicate, PredicateBasedState, State, @@ -202,6 +205,260 @@ def test_frozen(self): [self.ball_has_color_red, self.ball_has_shape_square] ) +# ------------------------- +# Serialization test domain +# ------------------------- + + +class Position(Enum): + LEFT = auto() + MIDDLE = auto() + RIGHT = auto() + + +class Velocity(Enum): + LEFT = auto() + RIGHT = auto() + + +class Angle(Enum): + STANDING = auto() + FALLING_LEFT = auto() + FALLING_RIGHT = auto() + + +class UnknownPredicateType(Enum): + UNKNOWN = auto() + + +class TestSerializableDiscretizer(Discretizer): + PREDICATE_TYPES = ( + Position, + Velocity, + Angle, + ) + + def discretize(self, non_discrete_state) -> State: + return non_discrete_state + + def state_to_str(self, state: State) -> str: + return super().state_to_str(state) + + def str_to_state(self, state_str: str) -> State: + return super().str_to_state(state_str) + + def nearest_state(self, state) -> Iterator[State]: + return iter(()) + + def all_actions(self) -> Sequence[Action]: + return [0, 1] + + def get_predicate_space(self) -> Sequence[Predicate]: + return [ + Predicate(Position.LEFT), + Predicate(Position.MIDDLE), + Predicate(Position.RIGHT), + Predicate(Velocity.LEFT), + Predicate(Velocity.RIGHT), + Predicate(Angle.STANDING), + Predicate(Angle.FALLING_LEFT), + Predicate(Angle.FALLING_RIGHT), + ] + + +class TestDiscretizerSerialization(unittest.TestCase): + def setUp(self): + self.discretizer = TestSerializableDiscretizer() + + self.position = Predicate(Position.MIDDLE) + self.velocity = Predicate(Velocity.RIGHT) + self.angle = Predicate(Angle.STANDING) + + # ------------------------- + # HAPPY PATHS + # ------------------------- + + def test_state_to_str_serializes_tuple_in_deterministic_order(self): + state = ( + self.angle, + self.position, + self.velocity, + ) + + state_str = self.discretizer.state_to_str(state) + + self.assertEqual( + state_str, + ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ), + ) + + def test_state_to_str_serializes_predicate_based_state_in_deterministic_order(self): + state = PredicateBasedState( + predicates={ + self.angle, + self.position, + self.velocity, + } + ) + + state_str = self.discretizer.state_to_str(state) + + self.assertEqual( + state_str, + ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ), + ) + + def test_str_to_state_handles_predicates_in_expected_order(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + state = self.discretizer.str_to_state(state_str) + + self.assertEqual( + state, + ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(Angle.STANDING), + ), + ) + + def test_str_to_state_handles_unordered_predicates(self): + state_str = ( + "Predicate(Angle.STANDING)&" + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)" + ) + + state = self.discretizer.str_to_state(state_str) + + self.assertEqual( + state, + ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(Angle.STANDING), + ), + ) + + def test_state_serialization_round_trip_with_unordered_tuple_state(self): + original_state = ( + Predicate(Angle.FALLING_LEFT), + Predicate(Velocity.LEFT), + Predicate(Position.RIGHT), + ) + + serialized = self.discretizer.state_to_str(original_state) + loaded_state = self.discretizer.str_to_state(serialized) + + self.assertEqual( + loaded_state, + ( + Predicate(Position.RIGHT), + Predicate(Velocity.LEFT), + Predicate(Angle.FALLING_LEFT), + ), + ) + + # ------------------------- + # UNHAPPY PATHS + # ------------------------- + + def test_str_to_state_raises_error_when_position_is_missing(self): + state_str = ( + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_velocity_is_missing(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_angle_is_missing(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_when_predicate_type_is_duplicated(self): + state_str = ( + "Predicate(Position.LEFT)&" + "Predicate(Position.RIGHT)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_for_unknown_enum_type(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(UnknownPredicateType.UNKNOWN)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_for_unknown_enum_member(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Predicate(Velocity.RIGHT)&" + "Predicate(Angle.UNKNOWN)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_str_to_state_raises_error_for_invalid_predicate_format(self): + state_str = ( + "Predicate(Position.MIDDLE)&" + "Velocity.RIGHT&" + "Predicate(Angle.STANDING)" + ) + + with self.assertRaises(ValueError): + self.discretizer.str_to_state(state_str) + + def test_state_to_str_raises_error_for_unknown_predicate_type(self): + state = ( + Predicate(Position.MIDDLE), + Predicate(Velocity.RIGHT), + Predicate(UnknownPredicateType.UNKNOWN), + ) + + with self.assertRaises(ValueError): + self.discretizer.state_to_str(state) + + def test_state_to_str_raises_error_for_unsupported_state_type(self): + unsupported_state = "not a valid state" + + with self.assertRaises(ValueError): + self.discretizer.state_to_str(unsupported_state) + if __name__ == "__main__": unittest.main()