WIP. (#12706)

rllib/agents/trainer_template.py

@@ -65,7 +65,7 @@ def build_trainer(
             Optional callable that takes the config to check for correctness.
             It may mutate the config as needed.
         default_policy (Optional[Type[Policy]]): The default Policy class to
-            use.
+            use if `get_policy_class` returns None.
         get_policy_class (Optional[Callable[
             TrainerConfigDict, Optional[Type[Policy]]]]): Optional callable
             that takes a config and returns the policy class or None. If None

rllib/evaluation/worker_set.py

@@ -79,7 +79,10 @@ class WorkerSet:
                 remote_spaces = ray.get(self.remote_workers(
                 )[0].foreach_policy.remote(
                     lambda p, pid: (pid, p.observation_space, p.action_space)))
-                spaces = {e[0]: (e[1], e[2]) for e in remote_spaces}
+                spaces = {
+                    e[0]: (getattr(e[1], "original_space", e[1]), e[2])
+                    for e in remote_spaces
+                }
             else:
                 spaces = None