To use your trained RLlib model for inference, you should load the RLModule from the checkpoint and preprocess observations exactly as during training, including any connector logic (e.g., normalization, concatenation). If you skip the connector pipeline or preprocess differently, inference results may be inconsistent or suboptimal. RLlib recommends using the EnvToModulePipeline and ModuleToEnvPipeline from the checkpoint for preprocessing and postprocessing, especially if you used custom connectors or filters during training. Directly calling .forward_inference() on the module with manually preprocessed observations can lead to mismatches and random behavior, as you observed.
For correct inference, load the connector pipelines from the checkpoint and use them to process observations and actions, ensuring consistency with training. See the official policy_inference_after_training_w_connector.py example for a full workflow. If you only use .forward_inference() directly, make sure your preprocessing exactly matches what the connector did during training, or else use the connector pipeline itself.
Would you like a step-by-step code example for this process?
Sources:
- policy_inference_after_training_w_connector.py
- policy_inference_after_training.py
- Ray RLlib inference discussion
- Ray RLlib inference GitHub issue
Hint: Mention @RunLLM in the post for followups.