Adds make_policy to PWIL and uses it in the baseline.
@@ -189,3 +189,9 @@ class PWILBuilder(builders.ActorLearnerBuilder[DirectRLNetworks,
) -> core.Actor:
return self._rl_agent.make_actor(random_key, policy, environment_spec,
variable_source, adder)
+
+ def make_policy(self,
+ networks: DirectRLNetworks,
+ environment_spec: specs.EnvironmentSpec,
+ evaluation: bool = False) -> DirectPolicyNetwork:
+ return self._rl_agent.make_policy(networks, environment_spec, evaluation)
