new architecture

problem/rl_env/brax_env.py

@@ -1,28 +1,13 @@
-from dataclasses import dataclass
-from typing import Callable
-
 import jax.numpy as jnp
 from brax import envs
-from core import State
-from .rl_jit import RLEnv, RLEnvConfig
 
-
-@dataclass(frozen=True)
-class BraxConfig(RLEnvConfig):
-    env_name: str = "ant"
-    backend: str = "generalized"
-
-    def __post_init__(self):
-        # TODO: Check if env_name is registered
-        # assert self.env_name in gymnax.registered_envs, f"Env {self.env_name} not registered"
-        pass
+from .rl_jit import RLEnv
 
 
 class BraxEnv(RLEnv):
-    def __init__(self, config: BraxConfig = BraxConfig()):
-        super().__init__(config)
-        self.config = config
-        self.env = envs.create(env_name=config.env_name, backend=config.backend)
+    def __init__(self, env_name: str = "ant", backend: str = "generalized"):
+        super().__init__()
+        self.env = envs.create(env_name=env_name, backend=backend)
 
     def env_step(self, randkey, env_state, action):
         state = self.env.step(env_state, action)
@@ -40,9 +25,7 @@ class BraxEnv(RLEnv):
     def output_shape(self):
         return (self.env.action_size,)
 
-    def show(self, randkey, state: State, act_func: Callable, params, save_path=None, height=512, width=512,
-             duration=0.1, *args,
-             **kwargs):
+    def show(self, randkey, state, act_func, params, save_path=None, height=512, width=512, duration=0.1, *args, **kwargs):
 
         import jax
         import imageio
@@ -56,8 +39,7 @@ class BraxEnv(RLEnv):
 
         def step(key, env_state, obs):
             key, _ = jax.random.split(key)
-            net_out = act_func(state, obs, params)
-            action = self.config.output_transform(net_out)
+            action = act_func(state, obs, params)
             obs, env_state, r, done, _ = self.step(randkey, env_state, action)
             return key, env_state, obs, r, done
 
@@ -72,7 +54,6 @@ class BraxEnv(RLEnv):
         def create_gif(image_list, gif_name, duration):
             with imageio.get_writer(gif_name, mode='I', duration=duration) as writer:
                 for image in image_list:
-                    # 确保图像的数据类型正确
                     formatted_image = np.array(image, dtype=np.uint8)
                     writer.append_data(formatted_image)
 

problem/rl_env/gymnax_env.py

@@ -1,26 +1,15 @@
-from dataclasses import dataclass
-from typing import Callable
-
 import gymnax
 
-from core import State
-from .rl_jit import RLEnv, RLEnvConfig
+from .rl_jit import RLEnv
 
 
-@dataclass(frozen=True)
-class GymNaxConfig(RLEnvConfig):
-    env_name: str = "CartPole-v1"
-
-    def __post_init__(self):
-        assert self.env_name in gymnax.registered_envs, f"Env {self.env_name} not registered"
-
 
 class GymNaxEnv(RLEnv):
 
-    def __init__(self, config: GymNaxConfig = GymNaxConfig()):
-        super().__init__(config)
-        self.config = config
-        self.env, self.env_params = gymnax.make(config.env_name)
+    def __init__(self, env_name):
+        super().__init__()
+        assert env_name in gymnax.registered_envs, f"Env {env_name} not registered"
+        self.env, self.env_params = gymnax.make(env_name)
 
     def env_step(self, randkey, env_state, action):
         return self.env.step(randkey, env_state, action, self.env_params)
@@ -36,5 +25,5 @@ class GymNaxEnv(RLEnv):
     def output_shape(self):
         return self.env.action_space(self.env_params).shape
 
-    def show(self, randkey, state: State, act_func: Callable, params):
+    def show(self, randkey, state, act_func, params, *args, **kwargs):
         raise NotImplementedError("GymNax render must rely on gym 0.19.0(old version).")

problem/rl_env/rl_jit.py

@@ -1,28 +1,18 @@
-from dataclasses import dataclass
-from typing import Callable
 from functools import partial
 
 import jax
 
-from config import ProblemConfig
+from .. import BaseProblem
 
-from core import Problem, State
-
-
-@dataclass(frozen=True)
-class RLEnvConfig(ProblemConfig):
-    output_transform: Callable = lambda x: x
-
-
-class RLEnv(Problem):
+class RLEnv(BaseProblem):
 
     jitable = True
 
-    def __init__(self, config: RLEnvConfig = RLEnvConfig()):
-        super().__init__(config)
-        self.config = config
+    # TODO: move output transform to algorithm
+    def __init__(self):
+        super().__init__()
 
-    def evaluate(self, randkey, state: State, act_func: Callable, params):
+    def evaluate(self, randkey, state, act_func, params):
         rng_reset, rng_episode = jax.random.split(randkey)
         init_obs, init_env_state = self.reset(rng_reset)
 
@@ -31,8 +21,7 @@ class RLEnv(Problem):
             return ~done
         def body_func(carry):
             obs, env_state, rng, _, tr = carry  # total reward
-            net_out = act_func(state, obs, params)
-            action = self.config.output_transform(net_out)
+            action = act_func(state, obs, params)
             next_obs, next_env_state, reward, done, _ = self.step(rng, env_state, action)
             next_rng, _ = jax.random.split(rng)
             return next_obs, next_env_state, next_rng, done, tr + reward
@@ -67,5 +56,5 @@ class RLEnv(Problem):
     def output_shape(self):
         raise NotImplementedError
 
-    def show(self, randkey, state: State, act_func: Callable, params, *args, **kwargs):
+    def show(self, randkey, state, act_func, params, *args, **kwargs):
         raise NotImplementedError