add more rl task in examples

2023-08-09 18:01:21 +08:00
parent af54db3b12
commit 3b6fe7eadc
18 changed files with 431 additions and 12 deletions
--- a/problem/rl_env/gymnax_env.py
+++ b/problem/rl_env/gymnax_env.py
@@ -1,12 +1,10 @@
 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
+from .rl_jit import RLEnv, RLEnvConfig


@dataclass(frozen=True)