Sort imports

fastfiz_env/__init__.py

@@ -26,9 +26,9 @@
 
 __version__ = "0.0.1"
 
-from .make import make, make_wrapped_env, make_callable_wrapped_env
-from .reward_functions import DefaultReward, RewardFunction, CombinedReward
-from . import envs, utils, wrappers, reward_functions
+from . import envs, reward_functions, utils, wrappers
+from .make import make, make_callable_wrapped_env, make_wrapped_env
+from .reward_functions import CombinedReward, DefaultReward, RewardFunction
 
 __all__ = [
     "make",
@@ -45,7 +45,6 @@
 
 from gymnasium.envs.registration import register
 
-
 register(
     id="FastFiz-v0",
     entry_point="fastfiz_env.envs:FastFiz",

fastfiz_env/envs/fastfiz.py

@@ -1,19 +1,20 @@
 from logging import warn
+from typing import Optional
+
 import fastfiz as ff
-import numpy as np
 import gymnasium as gym
+import numpy as np
 from gymnasium import spaces
 
 from fastfiz_env.utils.fastfiz.fastfiz import num_balls_pocketed
 
+from ..reward_functions import DefaultReward, RewardFunction
 from ..utils.fastfiz import (
     create_random_table_state,
     get_ball_positions,
     normalize_ball_positions,
 )
-from .utils import game_won, terminal_state, possible_shot
-from typing import Optional
-from ..reward_functions import RewardFunction, DefaultReward
+from .utils import game_won, possible_shot, terminal_state
 
 
 class FastFiz(gym.Env):

fastfiz_env/envs/pockets_fastfiz.py

@@ -1,21 +1,22 @@
 from logging import warn
+from typing import Optional
+
 import fastfiz as ff
-import numpy as np
 import gymnasium as gym
+import numpy as np
 from gymnasium import spaces
 
+from ..reward_functions import DefaultReward, RewardFunction
 from ..utils.fastfiz import (
+    ball_state_to_pocket,
     create_random_table_state,
     get_ball_positions,
-    normalize_ball_positions,
-    pocket_centers,
-    ball_state_to_pocket,
     get_pocket_center,
+    normalize_ball_positions,
     num_balls_pocketed,
+    pocket_centers,
 )
-from .utils import game_won, terminal_state, possible_shot
-from typing import Optional
-from ..reward_functions import RewardFunction, DefaultReward
+from .utils import game_won, possible_shot, terminal_state
 
 
 class PocketsFastFiz(gym.Env):

fastfiz_env/make.py

@@ -1,8 +1,9 @@
-from gymnasium.envs.registration import EnvSpec
 import gymnasium as gym
+from gymnasium.envs.registration import EnvSpec
 
 from fastfiz_env.wrappers.action import ActionSpaces, FastFizActionWrapper
-from .reward_functions import RewardFunction, DefaultReward
+
+from .reward_functions import DefaultReward, RewardFunction
 
 
 def make(

fastfiz_env/reward_functions/__init__.py

@@ -2,12 +2,12 @@
 This module contains the reward functions used in the FastFiz environment.
 """
 
-from .reward_function import RewardFunction, Weight
-from .combined_reward import CombinedReward
+from . import common
 from .binary_reward import BinaryReward
+from .combined_reward import CombinedReward
 from .default_reward import DefaultReward
+from .reward_function import RewardFunction, Weight
 from .winning_reward import WinningReward
-from . import common
 
 __all__ = [
     "RewardFunction",

fastfiz_env/reward_functions/binary_reward.py

@@ -1,9 +1,11 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Union
+
 import fastfiz as ff
-from .reward_function import RewardFunction, Weight
 import numpy as np
 
+from .reward_function import RewardFunction, Weight
+
 
 class BinaryReward(RewardFunction, ABC):
     def __init__(

fastfiz_env/reward_functions/combined_reward.py

@@ -1,9 +1,11 @@
 from typing import Optional
-from .reward_function import RewardFunction, Weight
-from .binary_reward import BinaryReward
+
 import fastfiz as ff
 import numpy as np
 
+from .binary_reward import BinaryReward
+from .reward_function import RewardFunction, Weight
+
 
 class CombinedReward(RewardFunction):
     """

fastfiz_env/reward_functions/common/__init__.py

@@ -2,33 +2,31 @@
 This module contains implementations of useful reward functions for the FastFiz environment.
 """
 
-from .step_pocketed_reward import StepPocketedReward
-from .total_distance_reward import TotalDistanceReward
-from .delta_best_total_distance_reward import DeltaBestTotalDistanceReward
-from .cue_ball_pocketed_reward import CueBallPocketedReward
+from .balls_not_moved_reward import BallsNotMovedReward
+from .constant_reward import ConstantReward
 from .cue_ball_not_moved_reward import CueBallNotMovedReward
+from .cue_ball_pocketed_reward import CueBallPocketedReward
+from .delta_best_total_distance_reward import DeltaBestTotalDistanceReward
+from .exponential_velocity_reward import ExponentialVelocityReward
 from .game_won_reward import GameWonReward
 from .impossible_shot_reward import ImpossibleShotReward
-from .constant_reward import ConstantReward
-from .balls_not_moved_reward import BallsNotMovedReward
-from .velocity_reward import VelocityReward
-from .exponential_velocity_reward import ExponentialVelocityReward
 from .step_no_balls_pocketed_reward import StepNoBallsPocketedReward
-
+from .step_pocketed_reward import StepPocketedReward
+from .total_distance_reward import TotalDistanceReward
+from .velocity_reward import VelocityReward
 from .weights import (
     ConstantWeight,
-    NegativeConstantWeight,
+    ConstantWeightBalls,
+    ConstantWeightCurrentStep,
     ConstantWeightMaxSteps,
-    NegativeConstantWeightMaxSteps,
     ConstantWeightNumBalls,
-    NegativeConstantWeightNumBalls,
-    ConstantWeightBalls,
+    NegativeConstantWeight,
     NegativeConstantWeightBalls,
-    ConstantWeightCurrentStep,
     NegativeConstantWeightCurrentStep,
+    NegativeConstantWeightMaxSteps,
+    NegativeConstantWeightNumBalls,
 )
 
-
 __all__ = [
     # Reward functions
     "StepPocketedReward",

fastfiz_env/reward_functions/common/balls_not_moved_reward.py

@@ -1,5 +1,5 @@
-from .. import BinaryReward
 from ...utils.fastfiz import any_ball_has_moved, get_ball_positions
+from .. import BinaryReward
 
 
 class BallsNotMovedReward(BinaryReward):

fastfiz_env/reward_functions/common/constant_reward.py

@@ -1,7 +1,8 @@
-from ..reward_function import RewardFunction
 import fastfiz as ff
 import numpy as np
 
+from ..reward_function import RewardFunction
+
 
 class ConstantReward(RewardFunction):
     """

fastfiz_env/reward_functions/common/cue_ball_not_moved_reward.py

@@ -1,7 +1,8 @@
-from .. import BinaryReward
 import fastfiz as ff
 import numpy as np
 
+from .. import BinaryReward
+
 
 class CueBallNotMovedReward(BinaryReward):
     """

fastfiz_env/reward_functions/common/cue_ball_pocketed_reward.py

@@ -1,7 +1,8 @@
-from .. import BinaryReward
 import fastfiz as ff
 import numpy as np
 
+from .. import BinaryReward
+
 
 class CueBallPocketedReward(BinaryReward):
     """

fastfiz_env/reward_functions/common/delta_best_total_distance_reward.py

@@ -1,12 +1,13 @@
-from ..reward_function import RewardFunction
+import fastfiz as ff
+import numpy as np
+
 from ...utils.fastfiz import (
     distances_to_closest_pocket,
     get_ball_positions,
     num_balls_in_play,
     pocket_centers,
 )
-import fastfiz as ff
-import numpy as np
+from ..reward_function import RewardFunction
 
 
 class DeltaBestTotalDistanceReward(RewardFunction):

fastfiz_env/reward_functions/common/exponential_velocity_reward.py

@@ -1,7 +1,8 @@
-from ..reward_function import RewardFunction
 import fastfiz as ff
 import numpy as np
 
+from ..reward_function import RewardFunction
+
 
 class ExponentialVelocityReward(RewardFunction):
     """

fastfiz_env/reward_functions/common/game_won_reward.py

@@ -1,7 +1,8 @@
-from .. import BinaryReward
 import fastfiz as ff
 import numpy as np
 
+from .. import BinaryReward
+
 
 class GameWonReward(BinaryReward):
     """

fastfiz_env/reward_functions/common/impossible_shot_reward.py

@@ -1,7 +1,8 @@
-from .. import BinaryReward
 import fastfiz as ff
 import numpy as np
 
+from .. import BinaryReward
+
 
 class ImpossibleShotReward(BinaryReward):
     """

fastfiz_env/reward_functions/common/step_no_balls_pocketed_reward.py

@@ -1,8 +1,9 @@
-from ...reward_functions import BinaryReward
-from ...utils.fastfiz.fastfiz import num_balls_pocketed
 import fastfiz as ff
 import numpy as np
 
+from ...reward_functions import BinaryReward
+from ...utils.fastfiz.fastfiz import num_balls_pocketed
+
 
 class StepNoBallsPocketedReward(BinaryReward):
     """

fastfiz_env/reward_functions/common/step_pocketed_reward.py

@@ -1,8 +1,9 @@
-from ..reward_function import RewardFunction
-from ...utils.fastfiz.fastfiz import num_balls_pocketed
 import fastfiz as ff
 import numpy as np
 
+from ...utils.fastfiz.fastfiz import num_balls_pocketed
+from ..reward_function import RewardFunction
+
 
 class StepPocketedReward(RewardFunction):
     """

fastfiz_env/reward_functions/common/total_distance_reward.py

@@ -1,12 +1,13 @@
+import fastfiz as ff
 import numpy as np
-from ..reward_function import RewardFunction
+
 from ...utils.fastfiz import (
     distances_to_closest_pocket,
     get_ball_positions,
     num_balls_in_play,
     pocket_centers,
 )
-import fastfiz as ff
+from ..reward_function import RewardFunction
 
 
 class TotalDistanceReward(RewardFunction):

fastfiz_env/reward_functions/common/velocity_reward.py

@@ -1,7 +1,8 @@
-from ..reward_function import RewardFunction
 import fastfiz as ff
 import numpy as np
 
+from ..reward_function import RewardFunction
+
 
 class VelocityReward(RewardFunction):
     """

fastfiz_env/reward_functions/default_reward.py

@@ -1,18 +1,17 @@
 from .combined_reward import CombinedReward
 from .common import (
+    # ExponentialVelocityReward,
+    BallsNotMovedReward,
     ConstantReward,
-    StepPocketedReward,
-    GameWonReward,
-    CueBallPocketedReward,
-    ConstantWeightBalls,
-    NegativeConstantWeightMaxSteps,
     ConstantWeight,
+    ConstantWeightBalls,
+    CueBallPocketedReward,
+    GameWonReward,
     NegativeConstantWeight,
-    # ExponentialVelocityReward,
-    BallsNotMovedReward,
+    NegativeConstantWeightMaxSteps,
+    StepPocketedReward,
 )
 
-
 rewards = [
     GameWonReward(ConstantWeight),
     CueBallPocketedReward(NegativeConstantWeight),

fastfiz_env/reward_functions/reward_function.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
+from typing import Callable, Optional, TypeAlias, Union
+
 import fastfiz as ff
-from typing import Union, Callable, Optional, TypeAlias
 import numpy as np
-from fastfiz_env.utils.fastfiz.fastfiz import num_balls_in_play
 
+from fastfiz_env.utils.fastfiz.fastfiz import num_balls_in_play
 
 Weight: TypeAlias = Union[float, Callable[[int, int, Optional[int]], float]]
 """

fastfiz_env/reward_functions/winning_reward.py

@@ -1,15 +1,15 @@
 from .combined_reward import CombinedReward
 from .common import (
     ConstantReward,
-    StepPocketedReward,
-    GameWonReward,
-    CueBallPocketedReward,
-    ConstantWeightBalls,
-    NegativeConstantWeightMaxSteps,
     ConstantWeight,
-    NegativeConstantWeight,
+    ConstantWeightBalls,
+    CueBallPocketedReward,
     ExponentialVelocityReward,
+    GameWonReward,
+    NegativeConstantWeight,
+    NegativeConstantWeightMaxSteps,
     StepNoBallsPocketedReward,
+    StepPocketedReward,
 )
 
 rewards = [