gr-libs 0.1.8__py3-none-any.whl → 0.2.2__py3-none-any.whl

gr_libs/environment/environment.py

@@ -1,14 +1,16 @@
+""" environment.py """
+
+import os
 from abc import abstractmethod
 from collections import namedtuple
-import os
 
 import gymnasium as gym
-from stable_baselines3.common.vec_env import DummyVecEnv
-from PIL import Image
 import numpy as np
 from gymnasium.envs.registration import register
-from minigrid.core.world_object import Wall, Lava
-from minigrid.wrappers import RGBImgPartialObsWrapper, ImgObsWrapper
+from minigrid.core.world_object import Lava, Wall
+from minigrid.wrappers import ImgObsWrapper, RGBImgPartialObsWrapper
+from PIL import Image
+from stable_baselines3.common.vec_env import DummyVecEnv
 
 MINIGRID, PANDA, PARKING, POINT_MAZE = "minigrid", "panda", "parking", "point_maze"
 
@@ -22,111 +24,207 @@ LSTMProperties = namedtuple(
 
 
 class EnvProperty:
+    """
+    Base class for environment properties.
+    """
+
     def __init__(self, name):
+        """
+        Initializes a new instance of the Environment class.
+
+        Args:
+            name (str): The name of the environment.
+        """
         self.name = name
 
     def __str__(self):
+        """
+        Returns a string representation of the object.
+        """
         return f"{self.name}"
 
     def __repr__(self):
+        """
+        Returns a string representation of the object.
+        """
         return f"{self.name}"
 
     def __eq__(self, other):
+        """
+        Check if this object is equal to another object.
+
+        Args:
+            other: The other object to compare with.
+
+        Returns:
+            True if the objects are equal, False otherwise.
+        """
         return self.name == other.name
 
     def __ne__(self, other):
+        """
+        Check if the current object is not equal to the other object.
+
+        Args:
+            other: The object to compare with.
+
+        Returns:
+            bool: True if the objects are not equal, False otherwise.
+        """
         return not self.__eq__(other)
 
     @abstractmethod
     def str_to_goal(self):
-        pass
+        """
+        Convert a problem name to a goal.
+        """
 
     @abstractmethod
     def gc_adaptable(self):
-        pass
+        """
+        Check if the environment is goal-conditioned adaptable.
+        """
 
     @abstractmethod
     def problem_list_to_str_tuple(self, problems):
-        pass
+        """
+        Convert a list of problems to a string tuple.
+        """
 
     @abstractmethod
     def goal_to_problem_str(self, goal):
-        pass
+        """
+        Convert a goal to a problem string.
+        """
 
     @abstractmethod
     def is_action_discrete(self):
-        pass
+        """
+        Check if the action space is discrete.
+        """
 
     @abstractmethod
     def is_state_discrete(self):
-        pass
+        """
+        Check if the state space is discrete.
+        """
 
     @abstractmethod
     def get_lstm_props(self):
-        pass
+        """
+        Get the LSTM properties for the environment.
+        """
 
     @abstractmethod
     def change_done_by_specific_desired(self, obs, desired, old_success_done):
-        pass
+        """
+        Change the 'done' flag based on a specific desired goal.
+        """
 
     @abstractmethod
     def is_done(self, done):
-        pass
+        """
+        Check if the episode is done.
+        """
 
     @abstractmethod
     def is_success(self, info):
-        pass
+        """
+        Check if the episode is successful.
+        """
 
     def create_vec_env(self, kwargs):
+        """
+        Create a vectorized environment.
+        """
         env = gym.make(**kwargs)
         return DummyVecEnv([lambda: env])
 
     @abstractmethod
     def change_goal_to_specific_desired(self, obs, desired):
-        pass
+        """
+        Change the goal to a specific desired goal.
+        """
 
 
 class GCEnvProperty(EnvProperty):
+    """
+    Base class for goal-conditioned environment properties.
+    """
+
     @abstractmethod
     def use_goal_directed_problem(self):
-        pass
+        """
+        Check if the environment uses a goal-directed problem.
+        """
 
     def problem_list_to_str_tuple(self, problems):
+        """
+        Convert a list of problems to a string tuple.
+        """
         return "goal_conditioned"
 
 
 class MinigridProperty(EnvProperty):
+    """
+    Environment properties for the Minigrid domain.
+    """
+
     def __init__(self, name):
         super().__init__(name)
         self.domain_name = "minigrid"
 
     def goal_to_problem_str(self, goal):
+        """
+        Convert a goal to a problem string.
+        """
         return self.name + f"-DynamicGoal-{goal[0]}x{goal[1]}-v0"
 
     def str_to_goal(self, problem_name):
+        """
+        Convert a problem name to a goal.
+        """
         parts = problem_name.split("-")
         goal_part = [part for part in parts if "x" in part]
         width, height = goal_part[0].split("x")
         return (int(width), int(height))
 
     def gc_adaptable(self):
+        """
+        Check if the environment is goal-conditioned adaptable.
+        """
         return False
 
     def problem_list_to_str_tuple(self, problems):
+        """
+        Convert a list of problems to a string tuple.
+        """
         return "_".join([f"[{s.split('-')[-2]}]" for s in problems])
 
     def is_action_discrete(self):
+        """
+        Check if the action space is discrete.
+        """
         return True
 
     def is_state_discrete(self):
+        """
+        Check if the state space is discrete.
+        """
         return True
 
     def get_lstm_props(self):
+        """
+        Get the LSTM properties for the environment.
+        """
         return LSTMProperties(
             batch_size=16, input_size=4, hidden_size=8, num_samples=40000
         )
 
     def create_sequence_image(self, sequence, img_path, problem_name):
+        """
+        Create a sequence image for the environment.
+        """
         if not os.path.exists(os.path.dirname(img_path)):
             os.makedirs(os.path.dirname(img_path))
         env_id = (
@@ -134,7 +232,7 @@ class MinigridProperty(EnvProperty):
             + "-DynamicGoal-"
             + problem_name.split("-DynamicGoal-")[1]
         )
-        result = register(
+        register(
             id=env_id,
             entry_point="gr_envs.minigrid_scripts.envs:CustomColorEnv",
             kwargs={
@@ -146,7 +244,6 @@ class MinigridProperty(EnvProperty):
                 "plan": sequence,
             },
         )
-        # print(result)
         env = gym.make(id=env_id)
         env = RGBImgPartialObsWrapper(env)  # Get pixel observations
         env = ImgObsWrapper(env)  # Get rid of the 'mission' field
@@ -156,34 +253,62 @@ class MinigridProperty(EnvProperty):
 
         ####### save image to file
         image_pil = Image.fromarray(np.uint8(img)).convert("RGB")
-        image_pil.save(r"{}.png".format(img_path))
+        image_pil.save(r"{}.png".format(os.path.join(img_path, "plan_image")))
 
     def change_done_by_specific_desired(self, obs, desired, old_success_done):
+        """
+        Change the 'done' flag based on a specific desired goal.
+        """
         assert (
             desired is None
         ), "In MinigridProperty, giving a specific 'desired' is not supported."
         return old_success_done
 
     def is_done(self, done):
+        """
+        Check if the episode is done.
+        """
         assert isinstance(done, np.ndarray)
         return done[0]
 
-    # Not used currently since TabularQLearner doesn't need is_success from the environment
     def is_success(self, info):
+        """
+        Check if the episode is successful.
+        """
         raise NotImplementedError("no other option for any of the environments.")
 
     def change_goal_to_specific_desired(self, obs, desired):
+        """
+        Change the goal to a specific desired goal.
+        """
         assert (
             desired is None
         ), "In MinigridProperty, giving a specific 'desired' is not supported."
 
 
 class PandaProperty(GCEnvProperty):
+    """
+    Environment properties for the Panda domain.
+    """
+
     def __init__(self, name):
+        """
+        Initialize a new instance of the Environment class.
+
+        Args:
+            name (str): The name of the environment.
+
+        Attributes:
+            domain_name (str): The domain name of the environment.
+
+        """
         super().__init__(name)
         self.domain_name = "panda"
 
     def str_to_goal(self, problem_name):
+        """
+        Convert a problem name to a goal.
+        """
         try:
             numeric_part = problem_name.split("PandaMyReachDenseX")[1]
             components = [
@@ -194,38 +319,62 @@ class PandaProperty(GCEnvProperty):
             for component in components:
                 floats.append(float(component))
             return np.array([floats], dtype=np.float32)
-        except Exception as e:
+        except Exception:
             return "general"
 
     def goal_to_problem_str(self, goal):
+        """
+        Convert a goal to a problem string.
+        """
         goal_str = "X".join(
             [str(float(g)).replace(".", "y").replace("-", "M") for g in goal[0]]
         )
         return f"PandaMyReachDenseX{goal_str}-v3"
 
     def gc_adaptable(self):
+        """
+        Check if the environment is goal-conditioned adaptable.
+        """
         return True
 
     def use_goal_directed_problem(self):
+        """
+        Check if the environment uses a goal-directed problem.
+        """
         return False
 
     def is_action_discrete(self):
+        """
+        Check if the action space is discrete.
+        """
         return False
 
     def is_state_discrete(self):
+        """
+        Check if the state space is discrete.
+        """
         return False
 
     def get_lstm_props(self):
+        """
+        Get the LSTM properties for the environment.
+        """
         return LSTMProperties(
             batch_size=32, input_size=9, hidden_size=8, num_samples=20000
         )
 
     def sample_goal():
+        """
+        Sample a random goal.
+        """
         goal_range_low = np.array([-0.40, -0.40, 0.10])
         goal_range_high = np.array([0.2, 0.2, 0.10])
         return np.random.uniform(goal_range_low, goal_range_high)
 
     def change_done_by_specific_desired(self, obs, desired, old_success_done):
+        """
+        Change the 'done' flag based on a specific desired goal.
+        """
         if desired is None:
             return old_success_done
         assert isinstance(
@@ -241,70 +390,134 @@ class PandaProperty(GCEnvProperty):
             return old_success_done
 
     def is_done(self, done):
+        """
+        Check if the episode is done.
+        """
         assert isinstance(done, np.ndarray)
         return done[0]
 
     def is_success(self, info):
+        """
+        Check if the episode is successful.
+        """
         assert "is_success" in info[0].keys()
         return info[0]["is_success"]
 
     def change_goal_to_specific_desired(self, obs, desired):
+        """
+        Change the goal to a specific desired goal.
+        """
         if desired is not None:
             obs["desired_goal"] = desired
 
 
 class ParkingProperty(GCEnvProperty):
+    """
+    Environment properties for the Parking domain.
+    """
 
     def __init__(self, name):
+        """
+        Initialize a new environment object.
+
+        Args:
+            name (str): The name of the environment.
+
+        Attributes:
+            domain_name (str): The domain name of the environment.
+
+        """
         super().__init__(name)
         self.domain_name = "parking"
 
     def goal_to_problem_str(self, goal):
+        """
+        Convert a goal to a problem string.
+        """
         return self.name.split("-v0")[0] + f"-GI-{goal}-v0"
 
     def gc_adaptable(self):
+        """
+        Check if the environment is goal-conditioned adaptable.
+        """
         return True
 
     def is_action_discrete(self):
+        """
+        Check if the action space is discrete.
+        """
         return False
 
     def is_state_discrete(self):
+        """
+        Check if the state space is discrete.
+        """
         return False
 
     def use_goal_directed_problem(self):
+        """
+        Check if the environment uses a goal-directed problem.
+        """
         return True
 
     def get_lstm_props(self):
+        """
+        Get the LSTM properties for the environment.
+        """
         return LSTMProperties(
             batch_size=32, input_size=8, hidden_size=8, num_samples=20000
         )
 
     def change_done_by_specific_desired(self, obs, desired, old_success_done):
+        """
+        Change the 'done' flag based on a specific desired goal.
+        """
         assert (
             desired is None
         ), "In ParkingProperty, giving a specific 'desired' is not supported."
         return old_success_done
 
     def is_done(self, done):
+        """
+        Check if the episode is done.
+        """
         assert isinstance(done, np.ndarray)
         return done[0]
 
     def is_success(self, info):
+        """
+        Check if the episode is successful.
+        """
         assert "is_success" in info[0].keys()
         return info[0]["is_success"]
 
     def change_goal_to_specific_desired(self, obs, desired):
+        """
+        Change the goal to a specific desired goal.
+        """
         assert (
             desired is None
         ), "In ParkingProperty, giving a specific 'desired' is not supported."
 
 
 class PointMazeProperty(EnvProperty):
+    """Environment properties for the Point Maze domain."""
+
     def __init__(self, name):
+        """
+        Initializes a new instance of the Environment class.
+
+        Args:
+            name (str): The name of the environment.
+
+        Attributes:
+            domain_name (str): The domain name of the environment.
+        """
         super().__init__(name)
         self.domain_name = "point_maze"
 
     def str_to_goal(self):
+        """Convert a problem name to a goal."""
         parts = self.name.split("-")
         # Find the part containing the goal size (usually after "DynamicGoal")
         sizes_parts = [part for part in parts if "x" in part]
@@ -314,40 +527,62 @@ class PointMazeProperty(EnvProperty):
         return (int(width), int(height))
 
     def gc_adaptable(self):
+        """Check if the environment is goal-conditioned adaptable."""
         return False
 
     def problem_list_to_str_tuple(self, problems):
+        """Convert a list of problems to a string tuple."""
         return "_".join([f"[{s.split('-')[-1]}]" for s in problems])
 
     def is_action_discrete(self):
+        """Check if the action space is discrete."""
         return False
 
     def is_state_discrete(self):
+        """Check if the state space is discrete."""
         return False
 
     def get_lstm_props(self):
+        """
+        Get the LSTM properties for the environment.
+        """
         return LSTMProperties(
             batch_size=32, input_size=6, hidden_size=8, num_samples=20000
         )
 
     def goal_to_problem_str(self, goal):
+        """
+        Convert a goal to a problem string.
+        """
         return self.name + f"-Goal-{goal[0]}x{goal[1]}"
 
     def change_done_by_specific_desired(self, obs, desired, old_success_done):
+        """
+        Change the 'done' flag based on a specific desired goal.
+        """
         assert (
             desired is None
         ), "In PointMazeProperty, giving a specific 'desired' is not supported."
         return old_success_done
 
     def is_done(self, done):
+        """
+        Check if the episode is done.
+        """
         assert isinstance(done, np.ndarray)
         return done[0]
 
     def is_success(self, info):
+        """
+        Check if the episode is successful.
+        """
         assert "success" in info[0].keys()
         return info[0]["success"]
 
     def change_goal_to_specific_desired(self, obs, desired):
+        """
+        Change the goal to a specific desired goal.
+        """
         assert (
             desired is None
         ), "In ParkingProperty, giving a specific 'desired' is not supported."

gr_libs/metrics/__init__.py

@@ -1,3 +1,5 @@
+""" metrics for GR algorithms """
+
 from .metrics import (
     mean_p_value,
     mean_wasserstein_distance,