code: def model_import_test(config):
      # Get the abs-path to use (bazel-friendly).
      rllib_dir = Path(__file__).parent.parent
      import_file = (
          str(rllib_dir) + "/tests/data/model_weights/weights.h5"
      )
  
      for fw in framework_iterator(config, ["tf", "torch"]):
          config.model["custom_model"] = (
              "keras_model" if fw != "torch" else "torch_model"
          )
  
          agent = config.build()
  
          def current_weight(agent):
              if fw == "tf":
                  return agent.get_weights()[DEFAULT_POLICY_ID][
                      "default_policy/value/kernel"
                  ][0]
              elif fw == "torch":
                  return float(
                      agent.get_weights()[DEFAULT_POLICY_ID][
                          "value_branch.weight"
                      ][0][0]
                  )
              else:
                  return agent.get_weights()[DEFAULT_POLICY_ID][4][0]
  
          # Import weights for our custom model from an h5 file.
          weight_before_import = current_weight(agent)
          agent.import_model(import_file=import_file)
          weight_after_import = current_weight(agent)
          check(weight_before_import, weight_after_import, false=True)
  
          # Train for a while.
          for _ in range(1):
              agent.train()
          weight_after_train = current_weight(agent)
          # Weights should have changed.
          check(weight_before_import, weight_after_train, false=True)
          check(weight_after_import, weight_after_train, false=True)
  
          # We can save the entire Agent and restore, weights should remain the
          # same.
          file = agent.save("after_train")
          check(weight_after_train, current_weight(agent))
          agent.restore(file)
          check(weight_after_train, current_weight(agent))
  
          # Import (untrained) weights again.
          agent.import_model(import_file=import_file)
          check(current_weight(agent), weight_after_import)
  
imports: import h5py
  import numpy
  from pathlib import Path
  import unittest
  import ray
  from ray.rllib.algorithms.ppo import PPOConfig
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.tf.misc import normc_initializer
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 1 for calling a global function with example usage: model_import_test(config) and returns: agent, float, agent