Relax some state constraints

Author: RunnersNum40

lerax/policy/actor_critic/mlp.py

@@ -41,9 +41,9 @@ class MLPActorCriticPolicy[
     action_model: MLP
     log_std: Float[Array, " action_size"]
 
-    def __init__(
+    def __init__[StateType: AbstractEnvLikeState](
         self,
-        env: AbstractEnvLike[AbstractEnvLikeState, ActType, ObsType],
+        env: AbstractEnvLike[StateType, ActType, ObsType],
         *,
         feature_size: int = 64,
         feature_width: int = 64,

lerax/wrapper/base_wrapper.py

@@ -22,7 +22,7 @@ class AbstractWrapper[
     WrapperStateType: AbstractEnvLikeState,
     WrapperActType,
     WrapperObsType,
-    StateType: AbstractEnvState,
+    StateType: AbstractEnvLikeState,
     ActType,
     ObsType,
 ](AbstractEnvLike[WrapperStateType, WrapperActType, WrapperObsType]):
@@ -31,7 +31,7 @@ class AbstractWrapper[
     env: eqx.AbstractVar[AbstractEnvLike[StateType, ActType, ObsType]]
 
     @property
-    def unwrapped(self) -> AbstractEnv[StateType, ActType, ObsType]:
+    def unwrapped(self) -> AbstractEnv:
         """Return the unwrapped environment"""
         return self.env.unwrapped
 

lerax/wrapper/misc.py

@@ -5,7 +5,7 @@
 from jax import numpy as jnp
 from jaxtyping import Array, ArrayLike, Bool, Float, Int, Key
 
-from lerax.env import AbstractEnvLike, AbstractEnvLikeState, AbstractEnvState
+from lerax.env import AbstractEnvLike, AbstractEnvLikeState
 from lerax.space import AbstractSpace
 
 from .base_wrapper import (
@@ -14,7 +14,7 @@
 )
 
 
-class Identity[StateType: AbstractEnvState, ActType, ObsType](
+class Identity[StateType: AbstractEnvLikeState, ActType, ObsType](
     AbstractWrapper[StateType, ActType, ObsType, StateType, ActType, ObsType]
 ):
     env: AbstractEnvLike[StateType, ActType, ObsType]
@@ -33,7 +33,7 @@ def step(
         return self.env.step(state, action, key=key)
 
 
-class EpisodeStatisticsState[StateType: AbstractEnvState](AbstractWrapperState):
+class EpisodeStatisticsState[StateType: AbstractEnvLikeState](AbstractWrapperState):
     env_state: StateType
 
     episode_length: Int[Array, ""]
@@ -76,7 +76,7 @@ def info(self) -> dict:
         }
 
 
-class EpisodeStatistics[StateType: AbstractEnvState, ActType, ObsType](
+class EpisodeStatistics[StateType: AbstractEnvLikeState, ActType, ObsType](
     AbstractWrapper[
         EpisodeStatisticsState[StateType], ActType, ObsType, StateType, ActType, ObsType
     ]
@@ -137,7 +137,7 @@ def __init__(self, step_count: Int[ArrayLike, ""], env_state: StateType):
         self.env_state = env_state
 
 
-class TimeLimit[StateType: AbstractEnvState, ActType, ObsType](
+class TimeLimit[StateType: AbstractEnvLikeState, ActType, ObsType](
     AbstractWrapper[
         TimeLimitState[StateType], ActType, ObsType, StateType, ActType, ObsType
     ]
@@ -193,7 +193,7 @@ def close(self):
         self.env.close()
 
 
-class AutoClose[StateType: AbstractEnvState, ActType, ObsType]():
+class AutoClose[StateType: AbstractEnvLikeState, ActType, ObsType]():
     """
     Closes the environment automatically when it is deleted.
 

lerax/wrapper/transform_action.py

@@ -6,15 +6,15 @@
 from jax import numpy as jnp
 from jaxtyping import Array, Float, Key
 
-from lerax.env import AbstractEnvLike, AbstractEnvState
+from lerax.env import AbstractEnvLike, AbstractEnvLikeState
 from lerax.space import AbstractSpace, Box
 
 from .base_wrapper import AbstractWrapper
 from .utils import rescale_box
 
 
 class AbstractPureTransformActionWrapper[
-    WrapperActType, StateType: AbstractEnvState, ActType, ObsType
+    WrapperActType, StateType: AbstractEnvLikeState, ActType, ObsType
 ](
     AbstractWrapper[
         StateType,
@@ -47,9 +47,9 @@ def close(self):
         self.env.close()
 
 
-class TransformAction[WrapperActType, StateType: AbstractEnvState, ActType, ObsType](
-    AbstractPureTransformActionWrapper[WrapperActType, StateType, ActType, ObsType]
-):
+class TransformAction[
+    WrapperActType, StateType: AbstractEnvLikeState, ActType, ObsType
+](AbstractPureTransformActionWrapper[WrapperActType, StateType, ActType, ObsType]):
     """Apply a function to the action before passing it to the environment"""
 
     env: AbstractEnvLike[StateType, ActType, ObsType]
@@ -67,7 +67,7 @@ def __init__(
         self.action_space = action_space
 
 
-class ClipAction[StateType: AbstractEnvState, ObsType](
+class ClipAction[StateType: AbstractEnvLikeState, ObsType](
     AbstractPureTransformActionWrapper[
         Float[Array, " ..."], StateType, Float[Array, " ..."], ObsType
     ],
@@ -98,7 +98,7 @@ def clip(action: Float[Array, " ..."]) -> Float[Array, " ..."]:
         self.action_space = action_space
 
 
-class RescaleAction[StateType: AbstractEnvState, ObsType](
+class RescaleAction[StateType: AbstractEnvLikeState, ObsType](
     AbstractPureTransformActionWrapper[
         Float[Array, " ..."], StateType, Float[Array, " ..."], ObsType
     ],

lerax/wrapper/transform_observation.py

@@ -7,15 +7,15 @@
 from jax import numpy as jnp
 from jaxtyping import Array, Float, Key
 
-from lerax.env import AbstractEnvLike, AbstractEnvState
+from lerax.env import AbstractEnvLike, AbstractEnvLikeState
 from lerax.space import AbstractSpace, Box
 
 from .base_wrapper import AbstractWrapper
 from .utils import rescale_box
 
 
 class AbstractPureObservationWrapper[
-    WrapperObsType, StateType: AbstractEnvState, ActType, ObsType
+    WrapperObsType, StateType: AbstractEnvLikeState, ActType, ObsType
 ](AbstractWrapper[StateType, ActType, WrapperObsType, StateType, ActType, ObsType]):
     """
     Apply a pure function to every observation that leaves the environment.
@@ -42,7 +42,7 @@ def close(self):
         self.env.close()
 
 
-class ClipObservation[StateType: AbstractEnvState](
+class ClipObservation[StateType: AbstractEnvLikeState](
     AbstractPureObservationWrapper[
         Float[Array, " ..."], StateType, Float[Array, " ..."], Float[Array, " ..."]
     ],
@@ -71,7 +71,7 @@ def __init__(self, env: AbstractEnvLike):
         self.observation_space = env.observation_space
 
 
-class RescaleObservation[StateType: AbstractEnvState](
+class RescaleObservation[StateType: AbstractEnvLikeState](
     AbstractPureObservationWrapper[
         Float[Array, " ..."], StateType, Float[Array, " ..."], Float[Array, " ..."]
     ],
@@ -101,7 +101,7 @@ def __init__(
         self.observation_space = new_box
 
 
-class FlattenObservation[StateType: AbstractEnvState, ObsType](
+class FlattenObservation[StateType: AbstractEnvLikeState, ObsType](
     AbstractPureObservationWrapper[
         Float[Array, " flat"], StateType, Float[Array, " ..."], ObsType
     ]

lerax/wrapper/transform_reward.py

@@ -7,14 +7,14 @@
 from jax import numpy as jnp
 from jaxtyping import Array, ArrayLike, Float, Key
 
-from lerax.env import AbstractEnvLike, AbstractEnvState
+from lerax.env import AbstractEnvLike, AbstractEnvLikeState
 
 from .base_wrapper import AbstractWrapper
 
 
-class AbstractPureTransformRewardWrapper[StateType: AbstractEnvState, ActType, ObsType](
-    AbstractWrapper[StateType, ActType, ObsType, StateType, ActType, ObsType]
-):
+class AbstractPureTransformRewardWrapper[
+    StateType: AbstractEnvLikeState, ActType, ObsType
+](AbstractWrapper[StateType, ActType, ObsType, StateType, ActType, ObsType]):
     """
     Apply a *pure* (stateless) function to every reward emitted by the wrapped
     environment.
@@ -39,11 +39,11 @@ def close(self):
         self.env.close()
 
 
-class ClipReward[StateType: AbstractEnvState, ActType, ObsType](
+class ClipReward[StateType: AbstractEnvLikeState, ActType, ObsType](
     AbstractPureTransformRewardWrapper[StateType, ActType, ObsType]
 ):
     """
-    Element-wise clip of rewards:  `reward ↦ clamp(min, max)`.
+    Element-wise clip of rewards:  `reward -> clamp(min, max)`.
     """
 
     env: AbstractEnvLike[StateType, ActType, ObsType]