remove create_func....

problem/func_fit/__init__.py

@@ -0,0 +1,3 @@
+from .func_fit import FuncFit, FuncFitConfig
+from .xor import XOR
+from .xor3d import XOR3d

problem/func_fit/func_fit.py

@@ -0,0 +1,69 @@
+from typing import Callable
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from config import ProblemConfig
+from core import Problem, State
+
+
+@dataclass(frozen=True)
+class FuncFitConfig(ProblemConfig):
+    error_method: str = 'mse'
+
+    def __post_init__(self):
+        assert self.error_method in {'mse', 'rmse', 'mae', 'mape'}
+
+
+class FuncFit(Problem):
+
+    def __init__(self, config: FuncFitConfig = FuncFitConfig()):
+        self.config = config
+        super().__init__(config)
+
+    def evaluate(self, randkey, state: State, act_func: Callable, params):
+
+        predict = act_func(state, self.inputs, params)
+
+        if self.config.error_method == 'mse':
+            loss = jnp.mean((predict - self.targets) ** 2)
+
+        elif self.config.error_method == 'rmse':
+            loss = jnp.sqrt(jnp.mean((predict - self.targets) ** 2))
+
+        elif self.config.error_method == 'mae':
+            loss = jnp.mean(jnp.abs(predict - self.targets))
+
+        elif self.config.error_method == 'mape':
+            loss = jnp.mean(jnp.abs((predict - self.targets) / self.targets))
+
+        else:
+            raise NotImplementedError
+
+        return -loss
+
+    def show(self, randkey, state: State, act_func: Callable, params):
+        predict = act_func(state, self.inputs, params)
+        inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
+        loss = -self.evaluate(randkey, state, act_func, params)
+        msg = ""
+        for i in range(inputs.shape[0]):
+            msg += f"input: {inputs[i]}, target: {target[i]}, predict: {predict[i]}\n"
+        msg += f"loss: {loss}\n"
+        print(msg)
+
+    @property
+    def inputs(self):
+        raise NotImplementedError
+
+    @property
+    def targets(self):
+        raise NotImplementedError
+
+    @property
+    def input_shape(self):
+        raise NotImplementedError
+
+    @property
+    def output_shape(self):
+        raise NotImplementedError

problem/func_fit/func_fitting.py

@@ -1,21 +0,0 @@
-from dataclasses import dataclass
-from typing import Callable
-
-from config import ProblemConfig
-from core import Problem, State
-
-
-@dataclass(frozen=True)
-class FuncFitConfig:
-    pass
-
-
-class FuncFit(Problem):
-    def __init__(self, config: ProblemConfig):
-        self.config = ProblemConfig
-
-    def setup(self, state=State()):
-        pass
-
-    def evaluate(self, state: State, act_func: Callable, params):
-        pass

problem/func_fit/xor.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+from .func_fit import FuncFit, FuncFitConfig
+
+
+class XOR(FuncFit):
+
+    def __init__(self, config: FuncFitConfig = FuncFitConfig()):
+        self.config = config
+        super().__init__(config)
+
+    @property
+    def inputs(self):
+        return np.array([
+            [0, 0],
+            [0, 1],
+            [1, 0],
+            [1, 1]
+        ])
+
+    @property
+    def targets(self):
+        return np.array([
+            [0],
+            [1],
+            [1],
+            [0]
+        ])
+
+    @property
+    def input_shape(self):
+        return (4, 2)
+
+    @property
+    def output_shape(self):
+        return (4, 1)

problem/func_fit/xor3d.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+from .func_fit import FuncFit, FuncFitConfig
+
+
+class XOR3d(FuncFit):
+
+    def __init__(self, config: FuncFitConfig = FuncFitConfig()):
+        self.config = config
+        super().__init__(config)
+
+    @property
+    def inputs(self):
+        return np.array([
+            [0, 0, 0],
+            [0, 0, 1],
+            [0, 1, 0],
+            [0, 1, 1],
+            [1, 0, 0],
+            [1, 0, 1],
+            [1, 1, 0],
+            [1, 1, 1],
+        ])
+
+    @property
+    def targets(self):
+        return np.array([
+            [0],
+            [1],
+            [1],
+            [0],
+            [1],
+            [0],
+            [0],
+            [1]
+        ])
+
+    @property
+    def input_shape(self):
+        return (8, 3)
+
+    @property
+    def output_shape(self):
+        return (8, 1)

problem/rl_env/__init__.py

@@ -0,0 +1 @@
+from .gymnax_env import GymNaxEnv, GymNaxConfig

problem/rl_env/gymnax_env.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+import gymnax
+
+from core import State
+from .rl_env import RLEnv, RLEnvConfig
+
+
+@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 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_env.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+from typing import Callable
+from functools import partial
+
+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):
+
+    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
+
+    @partial(jax.jit, static_argnums=(0,))
+    def step(self, randkey, env_state, action):
+        return self.env_step(randkey, env_state, action)
+
+    @partial(jax.jit, static_argnums=(0,))
+    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