Merge pull request #24 from YannBerthelot/gymnax_unification

changed from t to time to unify with gymnax

Author: YannBerthelot

src/target_gym/base.py

@@ -9,4 +9,4 @@ class EnvParams:
 
 @struct.dataclass
 class EnvState:
-    t: int
+    time: int

src/target_gym/benchmark_speed.py

@@ -47,12 +47,12 @@ def step_fn(carry, _):
 
     # Warmup
     state = step_fn_jit(state)
-    state.t.block_until_ready()
+    state.time.block_until_ready()
 
     # Timing
     t0 = time.time()
     state = step_fn_jit(state)
-    state.t.block_until_ready()
+    state.time.block_until_ready()
     t1 = time.time()
 
     total_steps = steps * batch_size

src/target_gym/bicycle/env.py

@@ -11,6 +11,7 @@
     jnp = None
 
 # Import your integrator
+from target_gym.base import EnvParams, EnvState
 from target_gym.integration import (
     integrate_dynamics,
 )
@@ -19,7 +20,7 @@
 
 
 @struct.dataclass
-class BikeState:
+class BikeState(EnvState):
     omega: float  # tilt angle [rad]
     omega_dot: float  # tilt angular velocity
     theta: float  # steering angle [rad]
@@ -30,7 +31,6 @@ class BikeState:
     x_b: float  # back wheel x
     y_b: float  # back wheel y
     last_d: float  # last displacement action (normalized)
-    t: int
     # for rendering
     torque: float = jnp.nan
     displacement: float = jnp.nan
@@ -41,7 +41,7 @@ def distance_from_start(self):
 
 
 @struct.dataclass
-class BikeParams:
+class BikeParams(EnvParams):
     c: float = 0.66
     dCM: float = 0.30
     h: float = 0.94
@@ -59,7 +59,6 @@ class BikeParams:
     delta_t: float = 0.05
 
     max_tilt_deg: float = 12.0
-    max_steps_in_episode: int = 1_000
 
     use_goal: bool = False
     goal_x: float = 0.0
@@ -229,7 +228,7 @@ def vecs_to_state(
         x_b=xb_new,
         y_b=yb_new,
         last_d=metrics.get("d", state.last_d).squeeze(),
-        t=state.t + 1,
+        time=state.time + 1,
     )
 
 
@@ -279,7 +278,7 @@ def _accel_fn(velocities, positions):
 def check_is_terminal(state: BikeState, params: BikeParams):
     max_tilt_rad = jnp.deg2rad(params.max_tilt_deg)
     terminated = jnp.abs(state.omega) > max_tilt_rad
-    truncated = state.t >= params.max_steps_in_episode
+    truncated = state.time >= params.max_steps_in_episode
     return terminated, truncated
 
 

src/target_gym/bicycle/env_jax.py

@@ -120,7 +120,7 @@ def reset_env(
             x_b=zero - params.l,
             y_b=zero,
             last_d=zero,
-            t=0,
+            time=0,
             torque=zero,
             displacement=zero,
         )

src/target_gym/bicycle/rendering.py

@@ -91,7 +91,7 @@ def world_to_screen(x, y):
     left_texts = [
         f"Lean ω: {np.rad2deg(state.omega):.1f}°",
         f"Steer θ: {np.rad2deg(state.theta):.1f}°",
-        f"Step: {state.t}",
+        f"Step: {state.time}",
     ]
     right_texts = [
         f"Heading ψ: {np.rad2deg(state.psi):.1f}°",

src/target_gym/car/env.py

@@ -125,7 +125,7 @@ def check_is_terminal(state: EnvState, params: CarParams, xp=jnp):
     terminated = jnp.logical_or(
         state.velocity <= params.min_velocity, state.velocity >= params.max_velocity
     )
-    truncated = state.t >= params.max_steps_in_episode
+    truncated = state.time >= params.max_steps_in_episode
     return terminated, truncated
 
 
@@ -223,7 +223,9 @@ def compute_next_state(
     )
 
     return (
-        state.replace(x=position, velocity=velocity, throttle=throttle, t=state.t + 1),
+        state.replace(
+            x=position, velocity=velocity, throttle=throttle, time=state.time + 1
+        ),
         metrics,
     )
 

src/target_gym/car/env_jax.py

@@ -104,7 +104,7 @@ def reset_env(
         )
 
         state = CarState(
-            t=0,
+            time=0,
             x=initial_x,
             velocity=initial_velocity,
             target_velocity=target_velocity,

src/target_gym/pc_gym/cstr/env.py

@@ -87,7 +87,7 @@ def compute_next_state(
         method=integration_method,
     )
     return (
-        state.replace(C_a=C_a, T=T, T_c=T_c_raw, t=state.t + 1),
+        state.replace(C_a=C_a, T=T, T_c=T_c_raw, time=state.time + 1),
         metrics,
     )
 
@@ -106,7 +106,7 @@ def check_is_terminal(state: CSTRState, params: CSTRParams, xp=jnp):
         state.C_a <= params.C_a_min, state.C_a >= params.C_a_max
     )
     terminated = jnp.logical_or(terminated_1, terminated_2)
-    truncated = state.t >= params.max_steps_in_episode
+    truncated = state.time >= params.max_steps_in_episode
     return terminated, truncated
 
 

src/target_gym/pc_gym/cstr/env_jax.py

@@ -104,7 +104,7 @@ def reset_env(
         )
 
         state = CSTRState(
-            t=0,
+            time=0,
             C_a=initial_C_a,
             T=params.initial_T,
             target_CA=initial_target_C_a,

src/target_gym/pc_gym/cstr/rendering.py

@@ -94,7 +94,7 @@ def _render(cls, screen, state, params, frames, clock, stride: int = 10):
     if not hasattr(cls, "history"):
         cls.history = {"t": [], "C_a": [], "T": [], "T_c": [], "reward": []}
 
-    step = state.t
+    step = state.time
     # Only render every `stride` steps
     if step % stride == 0 or step == 1:
         frame, cls.history = render_cstr(state, params, step, cls.history)