Export SAC in the Tensorflow format

RLlib
1

How severe does this issue affect your experience of using Ray?

  • High: It blocks me to complete my task.

I am currently trying to use RLlib to train the model then export it in the TensorFlow Saved Model format for further conversion. So I attempted to do it with custom Gymnasium Environment on modified example from the repository as shown below.

import ray
from env_simp_corr import SimpleCorridor
from ray.rllib.algorithms.sac import SACConfig
from ray.tune.registry import get_trainable_cls
import os

ray.init()

def train_and_export_policy_and_model(algo_name, num_steps, model_dir, ckpt_dir):
    # cls = get_trainable_cls(algo_name)
    # config = cls.get_default_config()
    config = SACConfig()
    # This Example is only for tf.
    config.framework("tf")
    config.export_native_model_files = True 
    config.environment(SimpleCorridor, env_config={"a": 0.2, "h": 300.0, "b": 0.3})
    alg = config.build()
    for _ in range(num_steps):
        print("Training...")
        alg.train()

    # Export Policy checkpoint.
    alg.export_policy_checkpoint(ckpt_dir)
    print("ckpt_dir: ", ckpt_dir)
    alg.export_policy_model(model_dir)
    # alg.export_model("CHECKPOINT", model_dir)
    print("model_dir: ", model_dir)
    

if __name__ == "__main__":
    algo = "SAC"
    model_dir = os.path.join(ray._private.utils.get_user_temp_dir(), "model_export_dir2")
    ckpt_dir = os.path.join(ray._private.utils.get_user_temp_dir(), "ckpt_export_dir2")
    num_steps = 3
    train_and_export_policy_and_model(algo, num_steps, model_dir, ckpt_dir)

and there are errors at the end of execution.

2023-07-07 22:30:10,426 WARNING tf_policy.py:649 -- Could not save keras model under self[TfPolicy].model.base_model!
    This is either due to ..
    a) .. this Policy's ModelV2 not having any `base_model` (tf.keras.Model) property
    b) .. the ModelV2's `base_model` not being used by the Algorithm and thus its
       variables not being properly initialized.

ckpt_dir:  /tmp/ckpt_export_dir2
2023-07-07 22:30:10,426 WARNING tf_policy.py:649 -- Could not save keras model under self[TfPolicy].model.base_model!
    This is either due to ..
    a) .. this Policy's ModelV2 not having any `base_model` (tf.keras.Model) property
    b) .. the ModelV2's `base_model` not being used by the Algorithm and thus its
       variables not being properly initialized.

So I would like to have some advice for the practical way to export the model that trained on the SAC algorithm for further use.

(I also did the same way in PPO algorithm, everything works fine and the TF Saved Model exported. but things changed when I change the algorithm to SAC.)