text: from rlberry.agents.torch import REINFORCEAgent
  from rlberry.envs.benchmarks.ball_exploration.ball2d import get_benchmark_env
  from rlberry.exploration_tools.discrete_counter import DiscreteCounter
  from rlberry.wrappers.uncertainty_estimator_wrapper import \
      UncertaintyEstimatorWrapper
  
  
  def test_reinforce_agent():
      _env = get_benchmark_env(level=1)
      n_episodes = 50
      horizon = 30
  
      #
      def uncertainty_estimator_fn(observation_space, action_space):
          counter = DiscreteCounter(observation_space, action_space,
                                    n_bins_obs=20)
          return counter
  
      env = UncertaintyEstimatorWrapper(_env,
                                        uncertainty_estimator_fn,
                                        bonus_scale_factor=1.0)
      #
      agent = REINFORCEAgent(env,
                             horizon=horizon,
                             gamma=0.99,
                             learning_rate=0.001,
                             use_bonus_if_available=True)
      agent.fit(budget=n_episodes)
      agent.policy(env.observation_space.sample())
  
  
  def test_reinforce_agent_partial_fit():
      env = get_benchmark_env(level=1)
      n_episodes = 10
      horizon = 30
  
      agent = REINFORCEAgent(env,
                             horizon=horizon,
                             gamma=0.99,
                             learning_rate=0.001,
                             use_bonus_if_available=False)
      agent.fit(budget=n_episodes // 2)
      agent.policy(env.observation_space.sample())
      assert agent.episode == 5
      agent.fit(budget=n_episodes // 2)
      assert agent.episode == 10
      agent.policy(env.observation_space.sample())
  