Add action_space_id option

P6-Pool · Apr 25, 2024 · ee54e63 · ee54e63
1 parent 26b4cf2
commit ee54e63
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/src/fastfiz_env/make.py b/src/fastfiz_env/make.py
@@ -43,6 +43,7 @@ def make_wrapped_env(
     num_balls: int,
     max_episode_steps: int,
     reward_function: RewardFunction,
+    action_space_id: ActionSpaces,
     **kwargs,
 ):
     """
@@ -56,15 +57,16 @@ def make_wrapped_env(
         disable_env_checker=False,
         **kwargs,
     )
-    env = FastFizActionWrapper(env, action_space_id=ActionSpaces.VECTOR_3D)
+    env = FastFizActionWrapper(env, action_space_id=action_space_id)
     return env
 
 
 def make_callable_wrapped_env(
     env_id: str,
     num_balls: int,
     max_episode_steps: int,
-    reward_function: RewardFunction,
+    reward_function: RewardFunction = DefaultReward,
+    action_space_id: ActionSpaces = ActionSpaces.VECTOR_3D,
     **kwargs,
 ):
     """
@@ -73,6 +75,6 @@ def make_callable_wrapped_env(
     """
 
     def _init() -> gym.Env:
-        return make_wrapped_env(env_id, num_balls, max_episode_steps, reward_function, **kwargs)
+        return make_wrapped_env(env_id, num_balls, max_episode_steps, reward_function, action_space_id, **kwargs)
 
     return _init
diff --git a/src/train.py b/src/train.py
@@ -13,6 +13,7 @@
     CallbackList,
 )
 from stable_baselines3.common.env_util import make_vec_env
+from fastfiz_env.wrappers.action import ActionSpaces
 from hyperparams import params_to_kwargs
 
 
@@ -40,11 +41,12 @@ def train(
     logs_path: str = "logs/",
     models_path: str = "models/",
     reward_function: RewardFunction = DefaultReward,
+    action_space_id: ActionSpaces = ActionSpaces.VECTOR_3D,
     callbacks=None,
     params: Optional[dict] = None,
 ) -> None:
     env = make_vec_env(
-        make_callable_wrapped_env(env_id, num_balls, max_episode_steps, reward_function),
+        make_callable_wrapped_env(env_id, num_balls, max_episode_steps, reward_function, action_space_id=action_space_id),
         n_envs=n_envs,
     )
 
@@ -122,6 +124,8 @@ def train(
         help="Path to hyperparameters file (file must have key 'params' with dict of hyperparameters",
     )
 
+    parser.add_argument("-a", "--action_id", type=ActionSpaces, choices=list(ActionSpaces), default=ActionSpaces.VECTOR_3D)
+
     args = parser.parse_args()
 
     reward_function = DefaultReward if args.reward == "DefaultReward" else WinningReward
@@ -164,5 +168,6 @@ def train(
         logs_path=logs_path,
         models_path=models_path,
         reward_function=reward_function,
+        action_space_id=args.action_id,
         params=params,
     )