def register_gym_env(class_entry_point, version="v0", kwargs=None):
  """Registers the class in Gym and returns the registered name and the env."""

  split_on_colon = class_entry_point.split(":")
  assert len(split_on_colon) == 2

  class_name = split_on_colon[1]
  # We have to add the version to conform to gym's API.
  env_name = "T2TEnv-{}-{}".format(class_name, version)
  gym.envs.register(id=env_name, entry_point=class_entry_point, kwargs=kwargs)

  logging.info(
      "Entry Point [%s] registered with id [%s]", class_entry_point, env_name)

  return env_name, gym.make(env_name)
Registers the class in Gym and returns the registered name and the env.
tensor2tensor/rl/gym_utils.py