Adds make_policy to AIL and uses it in the GAIL baseline.
@@ -322,3 +322,10 @@ class AILBuilder(builders.ActorLearnerBuilder[ail_networks.AILNetworks,
) -> core.Actor:
return self._rl_agent.make_actor(random_key, policy, environment_spec,
variable_source, adder)
+
+ def make_policy(self,
+ networks: ail_networks.AILNetworks,
+ environment_spec: specs.EnvironmentSpec,
+ evaluation: bool = False) -> DirectPolicyNetwork:
+ return self._rl_agent.make_policy(networks.direct_rl_networks,
+ environment_spec, evaluation)
