code: def test_identity_discrete(model_name):
      """
      Test if the algorithm (with a given policy)
      can learn an identity transformation (i.e. return observation as an action)
  
      :param model_name: (str) Name of the RL model
      """
      env = DummyVecEnv([lambda: IdentityEnv(10)])
  
      model = LEARN_FUNC_DICT[model_name](env)
      evaluate_policy(
          model, env, n_eval_episodes=20, reward_threshold=90
      )
  
      obs = env.reset()
      assert model.action_probability(obs).shape == (
          1,
          10,
      ), "Error: action_probability not returning correct shape"
      action = env.action_space.sample()
      action_prob = model.action_probability(obs, actions=action)
      assert (
          np.prod(action_prob.shape) == 1
      ), "Error: not scalar probability"
      action_logprob = model.action_probability(
          obs, actions=action, logp=True
      )
      assert np.allclose(action_prob, np.exp(action_logprob)), (
          action_prob,
          action_logprob,
      )
  
      # Free memory
      del model, env
  
file_path: stable-baselines/tests/test_identity.py
imports: import pytest
  import numpy
  from stable_baselines import A2C, ACER, ACKTR, DQN, DDPG, SAC, PPO1, PPO2, TD3, TRPO
  from stable_baselines.ddpg import NormalActionNoise
  from stable_baselines.common.identity_env import IdentityEnv, IdentityEnvBox
  from stable_baselines.common.vec_env import DummyVecEnv
  from stable_baselines.common.evaluation import evaluate_policy
label_desc: Use image node_id 1 for calling a global function with example usage: test_identity_discrete(model_name) without return types