classic-v1: continuous actions, and frames only on request

src/plugrl_env_client/envs/classic/classic_env.py

@@ -1,3 +1,15 @@
+"""Gymnasium's classic-control environments as a PlugRL environment.
+
+Both kinds of action space work. A discrete one (CartPole-v1) takes the
+action as an integer; a continuous one (Pendulum-v1, MountainCarContinuous-v0)
+as a float array of the space's shape and dtype. This env used to cast every
+action to an integer, so the continuous tasks could not run at all.
+
+Rendering is off by default, as in the MuJoCo family: a state-only policy
+never reads the frames, and rendering every step costs more than the physics.
+Turn it on with `--env.render` when something downstream wants pixels.
+"""
+
 import dataclasses
 import importlib.util
 import numpy as np
@@ -18,6 +30,7 @@
 @dataclasses.dataclass
 class ClassicConfig(BaseEnvConfig):
     name: str = "CartPole-v1"
+    render: bool = False
 
 
 @register_env(UID)
@@ -39,17 +52,22 @@ def __init__(
         )
         if self.num_envs != 1:
             raise ValueError("ClassicEnv only supports num_envs=1")
-        env = gym.make(config.name, render_mode="rgb_array")
+        self.render_frames = bool(config.render)
+        env = gym.make(
+            config.name, render_mode="rgb_array" if self.render_frames else None
+        )
         self.env = env
         self.game_name = config.name
+        self.discrete = isinstance(env.action_space, gym.spaces.Discrete)
         # rollout() sizes its action plan from this before the first step, so
         # an env without it cannot run at all.
         self.single_action_space = env.action_space
         self.action_space = env.action_space
 
     def prepare_obs(self, obs: np.ndarray) -> Observation:
-        frame = self.env.render()
-        frames = {"env": np.array(frame)[None, ...]}
+        frames = {}
+        if self.render_frames:
+            frames = {"env": np.array(self.env.render())[None, ...]}
         states = {"obs": obs[None, ...]}
         return Observation(
             images=frames,
@@ -66,7 +84,11 @@ def reset(
     def step(
         self, action: Action
     ) -> tuple[Observation | None, np.ndarray, np.ndarray, np.ndarray, dict]:
-        action = int(action.item())
+        if self.discrete:
+            action = int(np.asarray(action).item())
+        else:
+            space = self.env.action_space
+            action = np.asarray(action, dtype=space.dtype).reshape(space.shape)
         obs, reward, terminated, truncated, info = self.env.step(action)
         reward = np.array([float(reward)], dtype=np.float32)
         terminated = np.array([bool(terminated)], dtype=np.bool_)

tests/test_classic_env.py

@@ -0,0 +1,61 @@
+"""The classic-control family runs its continuous tasks, and renders only on request.
+
+It used to cast every action to an integer, which made Pendulum-v1 and
+MountainCarContinuous-v0 unusable, and it rendered a frame every step whether
+or not anything read it.
+"""
+
+import numpy as np
+import pytest
+
+classic_env = pytest.importorskip(
+    "plugrl_env_client.envs.classic.classic_env", exc_type=ImportError
+)
+
+ClassicConfig = classic_env.ClassicConfig
+ClassicEnv = classic_env.ClassicEnv
+
+
+def make(name: str, render: bool = False) -> ClassicEnv:
+    return ClassicEnv(config=ClassicConfig(name=name, render=render), num_envs=1)
+
+
+class TestContinuousActions:
+    def test_pendulum_takes_a_float_torque(self):
+        env = make("Pendulum-v1")
+        env.reset(seed=0)
+        env.env.unwrapped.state = np.array([0.0, 0.0])
+        # A torque of 1.5 from rest: an integer cast would have applied 1.0.
+        obs, reward, terminated, truncated, _ = env.step(
+            np.array([1.5], dtype=np.float32)
+        )
+        assert obs.states["obs"].shape == (1, 3)
+        assert obs.states["obs"][0, 2] == pytest.approx(3.0 * 1.5 * 0.05, rel=1e-6)
+        assert reward.shape == terminated.shape == truncated.shape == (1,)
+
+    def test_the_action_takes_the_space_shape(self):
+        env = make("Pendulum-v1")
+        env.reset(seed=0)
+        # rollout() hands over one env's row of the plan, possibly with extra
+        # leading axes; the env must reduce it to the space's own shape.
+        env.step(np.array([[0.5]], dtype=np.float32))
+
+
+class TestDiscreteActions:
+    def test_cartpole_still_takes_an_integer(self):
+        env = make("CartPole-v1")
+        env.reset(seed=0)
+        obs, *_ = env.step(np.array([1]))
+        assert obs.states["obs"].shape == (1, 4)
+
+
+class TestRendering:
+    def test_no_frames_by_default(self):
+        env = make("Pendulum-v1")
+        obs, _ = env.reset(seed=0)
+        assert obs.images == {}
+
+    def test_frames_when_asked(self):
+        env = make("Pendulum-v1", render=True)
+        obs, _ = env.reset(seed=0)
+        assert obs.images["env"].ndim == 4