diff --git a/bindings/pydrake/gym/_drake_gym_env.py b/bindings/pydrake/gym/_drake_gym_env.py index b6187615b66e..8a335c94a426 100644 --- a/bindings/pydrake/gym/_drake_gym_env.py +++ b/bindings/pydrake/gym/_drake_gym_env.py @@ -33,8 +33,8 @@ def __init__( self, simulator: Simulator | Callable[[RandomGenerator], Simulator], time_step: float, - action_space: gym.spaces.Space, - observation_space: gym.spaces.Space, + action_space: gym.spaces.Space | None, + observation_space: gym.spaces.Space | None, reward: Callable[[System, Context], float] | OutputPortIndex | str, action_port_id: InputPort | InputPortIndex | str = None, observation_port_id: OutputPortIndex | str = None, @@ -127,10 +127,12 @@ def __init__( assert time_step > 0 self.time_step = time_step - assert isinstance(action_space, gym.spaces.Space) + if action_space is not None: + assert isinstance(action_space, gym.spaces.Space) self.action_space = action_space - assert isinstance(observation_space, gym.spaces.Space) + if observation_space is not None: + assert isinstance(observation_space, gym.spaces.Space) self.observation_space = observation_space if isinstance(reward, (OutputPortIndex, str)): @@ -193,9 +195,21 @@ def _setup(self): else: self.action_port = system.GetInputPort(self.action_port_id) if self.action_port.get_data_type() == PortDataType.kVectorValued: + if self.action_space is None: + self.action_space = gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.action_port.size(),), + dtype=np.float64, + ) assert np.array_equal( self.action_space.shape, [self.action_port.size()] ) + else: + assert self.action_space is not None, ( + "action_space must be provided when the action port is not " + "vector-valued" + ) def get_output_port(id): if isinstance(id, OutputPortIndex): @@ -206,9 +220,21 @@ def get_output_port(id): if self.observation_port_id: self.observation_port = get_output_port(self.observation_port_id) if self.observation_port.get_data_type() == PortDataType.kVectorValued: + if self.observation_space is None: + self.observation_space = gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.observation_port.size(),), + dtype=np.float64, + ) assert np.array_equal( self.observation_space.shape, [self.observation_port.size()] ) + else: + assert self.observation_space is not None, ( + "observation_space must be provided when the observation " + "port is not vector-valued" + ) # Note: We require that there is no direct feedthrough action_port to # observation_port. Unfortunately, HasDirectFeedthrough returns false diff --git a/bindings/pydrake/gym/test/drake_gym_test.py b/bindings/pydrake/gym/test/drake_gym_test.py index 507de50e1049..3aaf4de178df 100644 --- a/bindings/pydrake/gym/test/drake_gym_test.py +++ b/bindings/pydrake/gym/test/drake_gym_test.py @@ -1,8 +1,15 @@ import unittest import gymnasium as gym +import numpy as np import stable_baselines3.common.env_checker +from pydrake.common.value import Value +from pydrake.gym import DrakeGymEnv +from pydrake.systems.analysis import Simulator +from pydrake.systems.framework import DiagramBuilder +from pydrake.systems.primitives import PassThrough + class DrakeGymTest(unittest.TestCase): """ @@ -65,3 +72,70 @@ def test_step(self): dut.reset() observation, _, _, _, _ = dut.step(dut.action_space.sample()) self.assertTrue(dut.observation_space.contains(observation)) + + def test_none_spaces_default_to_infinite_box(self): + """Passing None for vector-valued ports builds ±inf Boxes.""" + size = 3 + builder = DiagramBuilder() + plant = builder.AddSystem(PassThrough(vector_size=size)) + builder.ExportInput(plant.get_input_port(), "actions") + builder.ExportOutput(plant.get_output_port(), "observations") + diagram = builder.Build() + simulator = Simulator(diagram) + + dut = DrakeGymEnv( + simulator=simulator, + time_step=0.1, + action_space=None, + observation_space=None, + reward=lambda system, context: 0.0, + action_port_id="actions", + observation_port_id="observations", + ) + + self.assertIsInstance(dut.action_space, gym.spaces.Box) + self.assertEqual(dut.action_space.shape, (size,)) + self.assertTrue(np.all(np.isneginf(dut.action_space.low))) + self.assertTrue(np.all(np.isposinf(dut.action_space.high))) + + self.assertIsInstance(dut.observation_space, gym.spaces.Box) + self.assertEqual(dut.observation_space.shape, (size,)) + self.assertTrue(np.all(np.isneginf(dut.observation_space.low))) + self.assertTrue(np.all(np.isposinf(dut.observation_space.high))) + + def test_none_spaces_reject_non_vector_ports(self): + """None is only valid for vector-valued ports.""" + builder = DiagramBuilder() + system = builder.AddSystem( + PassThrough(abstract_model_value=Value("model")) + ) + builder.ExportInput(system.get_input_port(), "actions") + builder.ExportOutput(system.get_output_port(), "observations") + diagram = builder.Build() + simulator = Simulator(diagram) + # Any concrete Space is fine here; we only exercise the None path. + dummy_space = gym.spaces.Discrete(1) + + with self.assertRaisesRegex(AssertionError, "action_space must be"): + DrakeGymEnv( + simulator=simulator, + time_step=0.1, + action_space=None, + observation_space=dummy_space, + reward=lambda system, context: 0.0, + action_port_id="actions", + observation_port_id="observations", + ) + + with self.assertRaisesRegex( + AssertionError, "observation_space must be" + ): + DrakeGymEnv( + simulator=simulator, + time_step=0.1, + action_space=dummy_space, + observation_space=None, + reward=lambda system, context: 0.0, + action_port_id="actions", + observation_port_id="observations", + )