[RLLib] Change connector configuration callbacks to transform ConnectorPipelineV2
Description
Change RLlib’s connector configuration options to accept a callable that takes as an argument a ConnectorPipelineV2 and returns a ConnectorPipelineV2 with any user changes.
RLlib should construct the complete pipeline, including default connectors (if default connectors are enabled), pass it to the callback, and use the returned pipeline.
Something like...
def configure_pipeline(
pipeline: ConnectorPipelineV2,
) -> ConnectorPipelineV2:
pipeline.insert_after(
AddObservationsFromEpisodesToBatch,
MyCustomConnector(),
)
return pipeline
config.env_runners(
env_to_module_connector=configure_pipeline,
)The callback could mutate the provided pipeline or return a replacement. RLlib should validate that the result is a ConnectorPipelineV2.
Use case
The intrinsic curiosity module requires two connectors at specific positions in the learner pipeline.
The current API only adds custom connectors before the default connectors. As a result, the example must define custom DQNTorchLearner and PPOTorchLearner subclasses and modify the private _learner_connector attribute after Learner.build().
I think this is a weakness of the current connector configuration that could be addressed with this change.
Source: ray-project/ray