Merge pull request #110 from sony/feature/20231012-enable-tuple-actio…

…n-support

Feature/20231012 enable tuple action support

nnabla_rl/algorithms/a2c.py

@@ -366,7 +366,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_continuous_action_env()
+        return not env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/atrpo.py

@@ -405,7 +405,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/bcq.py

@@ -341,7 +341,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/bear.py

@@ -401,7 +401,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/categorical_dqn.py

@@ -339,7 +339,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_continuous_action_env()
+        return not env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()
 
     @classmethod
     def is_rnn_supported(self):

nnabla_rl/algorithms/common_utils.py

@@ -374,9 +374,13 @@ def __call__(self, s: Union[np.ndarray, Tuple[np.ndarray, ...]], *, begin_of_epi
             for key in self._rnn_internal_states.keys():
                 # copy internal states of previous iteration
                 self._rnn_internal_states[key].d = prev_rnn_states[key].d
-        self._action.forward(clear_no_need_grad=True)
-        # No need to save internal states
-        action = np.squeeze(self._action.d, axis=0) if batch_size == 1 else self._action.d
+        if self._env_info.is_tuple_action_env():
+            nn.forward_all(self._action, clear_no_need_grad=True)
+            action = tuple(np.squeeze(a.d, axis=0) if batch_size == 1 else a.d for a in self._action)
+        else:
+            self._action.forward(clear_no_need_grad=True)
+            # No need to save internal states
+            action = np.squeeze(self._action.d, axis=0) if batch_size == 1 else self._action.d
         return action, {}
 
     @abstractmethod

nnabla_rl/algorithms/ddp.py

@@ -308,7 +308,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def trainers(self):

nnabla_rl/algorithms/ddpg.py

@@ -359,7 +359,7 @@ def is_rnn_supported(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/decision_transformer.py

@@ -263,7 +263,8 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(
             env_or_env_info, gym.Env) else env_or_env_info
-        return env_info.is_continuous_action_env() or env_info.is_discrete_action_env()
+        return ((env_info.is_continuous_action_env() or env_info.is_discrete_action_env())
+                and not env_info.is_tuple_action_env())
 
     @classmethod
     def is_rnn_supported(self):

nnabla_rl/algorithms/demme_sac.py

@@ -624,7 +624,7 @@ def is_rnn_supported(cls):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/dqn.py

@@ -335,7 +335,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_continuous_action_env()
+        return not env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()
 
     @classmethod
     def is_rnn_supported(self):

nnabla_rl/algorithms/gail.py

@@ -543,7 +543,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/her.py

@@ -352,4 +352,4 @@ def is_supported_env(cls, env_or_env_info):
         # continuous action env
         is_continuous_action_env = env_info.is_continuous_action_env()
         is_goal_conditioned_env = env_info.is_goal_conditioned_env()
-        return (is_continuous_action_env and is_goal_conditioned_env)
+        return (is_continuous_action_env and is_goal_conditioned_env) and not env_info.is_tuple_action_env()

nnabla_rl/algorithms/icml2015_trpo.py

@@ -279,7 +279,9 @@ def _solvers(self):
 
     @classmethod
     def is_supported_env(cls, env_or_env_info):
-        return True  # supports all enviroments
+        env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
+            else env_or_env_info
+        return not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/icml2018_sac.py

@@ -442,7 +442,7 @@ def is_rnn_supported(cls):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/icra2018_qtopt.py

@@ -264,4 +264,4 @@ def _random_action_selector(self, s, *, begin_of_episode=False):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return env_info.is_continuous_action_env()
+        return env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()

nnabla_rl/algorithms/iqn.py

@@ -352,7 +352,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_continuous_action_env()
+        return not env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()
 
     @classmethod
     def is_rnn_supported(self):

nnabla_rl/algorithms/lqr.py

@@ -159,7 +159,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def trainers(self):

nnabla_rl/algorithms/mppi.py

@@ -365,7 +365,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/ppo.py

@@ -448,7 +448,9 @@ def _solvers(self):
 
     @classmethod
     def is_supported_env(cls, env_or_env_info):
-        return True  # supports all enviroments
+        env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
+            else env_or_env_info
+        return not env_info.is_tuple_action_env()
 
     def _build_ppo_actors(self, env, v_function, policy, state_preprocessor):
         actors = []

nnabla_rl/algorithms/qrdqn.py

@@ -325,7 +325,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_continuous_action_env()
+        return not env_info.is_continuous_action_env() and not env_info.is_tuple_action_env()
 
     @classmethod
     def is_rnn_supported(self):

nnabla_rl/algorithms/qrsac.py

@@ -395,7 +395,7 @@ def is_rnn_supported(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/reinforce.py

@@ -266,7 +266,9 @@ def _solvers(self):
 
     @classmethod
     def is_supported_env(cls, env_or_env_info):
-        return True  # supports all enviroments
+        env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
+            else env_or_env_info
+        return not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/sac.py

@@ -397,7 +397,7 @@ def is_rnn_supported(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/td3.py

@@ -384,7 +384,7 @@ def is_rnn_supported(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/trpo.py

@@ -390,7 +390,7 @@ def _solvers(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/algorithms/xql.py

@@ -420,7 +420,7 @@ def is_rnn_supported(self):
     def is_supported_env(cls, env_or_env_info):
         env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
             else env_or_env_info
-        return not env_info.is_discrete_action_env()
+        return not env_info.is_discrete_action_env() and not env_info.is_tuple_action_env()
 
     @property
     def latest_iteration_state(self):

nnabla_rl/environment_explorer.py

@@ -122,12 +122,7 @@ def rollout(self, env: gym.Env) -> List[Experience]:
     def _step_once(self, env, *, begin_of_episode=False) -> Tuple[Experience, bool]:
         self._steps += 1
         if self._steps < self._config.warmup_random_steps:
-            action_info: Dict[str, Any] = {}
-            if self._env_info.is_discrete_action_env():
-                action = env.action_space.sample()
-                self._action = np.asarray(action).reshape((1, ))
-            else:
-                self._action = env.action_space.sample()
+            self._action, action_info = self._warmup_action(env)
         else:
             self._action, action_info = self.action(self._steps,
                                                     cast(np.ndarray, self._state),
@@ -156,9 +151,30 @@ def _step_once(self, env, *, begin_of_episode=False) -> Tuple[Experience, bool]:
             self._state = self._next_state
         return experience, done
 
+    def _warmup_action(self, env):
+        return _sample_action(env, self._env_info)
+
 
 def _is_end_of_episode(done, timelimit, timelimit_as_terminal):
     if not done:
         return False
     else:
         return (not timelimit) or (timelimit and timelimit_as_terminal)
+
+
+def _sample_action(env, env_info):
+    action_info: Dict[str, Any] = {}
+    if env_info.is_tuple_action_env():
+        action = []
+        for a, action_space in zip(env.action_space.sample(), env_info.action_space):
+            if isinstance(action_space, gym.spaces.Discrete):
+                a = np.asarray(a).reshape((1, ))
+            action.append(a)
+        action = tuple(action)
+    else:
+        if env_info.is_discrete_action_env():
+            action = env.action_space.sample()
+            action = np.asarray(action).reshape((1, ))
+        else:
+            action = env.action_space.sample()
+    return action, action_info

nnabla_rl/environments/__init__.py

@@ -1,5 +1,5 @@
 # Copyright 2020,2021 Sony Corporation.
-# Copyright 2021,2022 Sony Group Corporation.
+# Copyright 2021,2022,2023 Sony Group Corporation.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,8 +17,10 @@
 
 from nnabla_rl.environments.dummy import (DummyAtariEnv, DummyContinuous, DummyContinuousActionGoalEnv, DummyDiscrete,  # noqa
                                           DummyDiscreteActionGoalEnv, DummyDiscreteImg, DummyContinuousImg,
-                                          DummyFactoredContinuous, DummyMujocoEnv, DummyTupleContinuous,
-                                          DummyTupleDiscrete, DummyTupleMixed)
+                                          DummyFactoredContinuous, DummyMujocoEnv,
+                                          DummyTupleContinuous, DummyTupleDiscrete, DummyTupleMixed,
+                                          DummyTupleStateContinuous, DummyTupleStateDiscrete,
+                                          DummyTupleActionContinuous, DummyTupleActionDiscrete)
 
 register(
     id='FakeMujocoNNablaRL-v1',