delete useless;

append readme

problem/rl_env/gym_env.py

@@ -1,40 +0,0 @@
-from dataclasses import dataclass
-from typing import Callable
-
-import gym
-
-from core import State
-from .rl_unjit import RLEnv, RLEnvConfig
-
-
-@dataclass(frozen=True)
-class GymConfig(RLEnvConfig):
-    env_name: str = "CartPole-v1"
-
-    def __post_init__(self):
-        assert self.env_name in gym.registered_envs, f"Env {self.env_name} not registered"
-
-
-class GymNaxEnv(RLEnv):
-
-    def __init__(self, config: GymConfig = GymConfig()):
-        super().__init__(config)
-        self.config = config
-        self.env, self.env_params = gym.make(config.env_name)
-
-    def env_step(self, randkey, env_state, action):
-        return self.env.step(randkey, env_state, action, self.env_params)
-
-    def env_reset(self, randkey):
-        return self.env.reset(randkey, self.env_params)
-
-    @property
-    def input_shape(self):
-        return self.env.observation_space(self.env_params).shape
-
-    @property
-    def output_shape(self):
-        return self.env.action_space(self.env_params).shape
-
-    def show(self, randkey, state: State, act_func: Callable, params):
-        raise NotImplementedError("GymNax render must rely on gym 0.19.0(old version).")

problem/rl_env/rl_unjit.py

@@ -1,69 +0,0 @@
-from dataclasses import dataclass
-from typing import Callable
-
-import jax
-
-from config import ProblemConfig
-
-from core import Problem, State
-
-
-@dataclass(frozen=True)
-class RLEnvConfig(ProblemConfig):
-    output_transform: Callable = lambda x: x
-
-
-class RLEnv(Problem):
-
-    jitable = False
-
-    def __init__(self, config: RLEnvConfig = RLEnvConfig()):
-        super().__init__(config)
-        self.config = config
-
-    def evaluate(self, randkey, state: State, act_func: Callable, params):
-        rng_reset, rng_episode = jax.random.split(randkey)
-        init_obs, init_env_state = self.reset(rng_reset)
-
-        def cond_func(carry):
-            _, _, _, done, _ = carry
-            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)
-            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
-
-        _, _, _, _, total_reward = jax.lax.while_loop(
-            cond_func,
-            body_func,
-            (init_obs, init_env_state, rng_episode, False, 0.0)
-        )
-
-        return total_reward
-
-    def step(self, randkey, env_state, action):
-        return self.env_step(randkey, env_state, action)
-
-    def reset(self, randkey):
-        return self.env_reset(randkey)
-
-    def env_step(self, randkey, env_state, action):
-        raise NotImplementedError
-
-    def env_reset(self, randkey):
-        raise NotImplementedError
-
-    @property
-    def input_shape(self):
-        raise NotImplementedError
-
-    @property
-    def output_shape(self):
-        raise NotImplementedError
-
-    def show(self, randkey, state: State, act_func: Callable, params):
-        raise NotImplementedError