def env_problem(env_problem_name, **kwargs):
  """Get and initialize the `EnvProblem` with the given name and batch size.

  Args:
    env_problem_name: string name of the registered env problem.
    **kwargs: forwarded to env problem's initialize method.

  Returns:
    an initialized EnvProblem with the given batch size.
  """

  ep_cls = Registries.env_problems[env_problem_name]
  ep = ep_cls()
  ep.initialize(**kwargs)
  return ep
Get and initialize the `EnvProblem` with the given name and batch size.

  Args:
    env_problem_name: string name of the registered env problem.
    **kwargs: forwarded to env problem's initialize method.

  Returns:
    an initialized EnvProblem with the given batch size.