diff --git a/examples/rl/rllib/rllib.py b/examples/rl/rllib/rllib.py
index c558478ed9..9260599c83 100644
--- a/examples/rl/rllib/rllib.py
+++ b/examples/rl/rllib/rllib.py
@@ -4,7 +4,7 @@
 import random
 from pathlib import Path
 from pprint import pprint as print
-from typing import Dict, Literal, Union
+from typing import Dict, Literal, Optional, Union
 
 import numpy as np
 
@@ -109,7 +109,7 @@ def main(
     resume_training,
     result_dir,
     checkpoint_freq: int,
-    checkpoint_num: int,
+    checkpoint_num: Optional[int],
     log_level: Literal["DEBUG", "INFO", "WARN", "ERROR"],
     save_model_path,
 ):