72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
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):
|
|
|
|
jitable = True
|
|
|
|
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, *args, **kwargs):
|
|
raise NotImplementedError
|