def test_discrete_env(rate_power): env = GridWorld() counter = DiscreteCounter(env.observation_space, env.action_space, rate_power=rate_power) for N in range(10, 20): assert counter.get_n_visited_states() == 0 assert counter.get_entropy() == 0.0 for ss in range(env.observation_space.n): for aa in range(env.action_space.n): for _ in range(N): ns, rr, _, _ = env.sample(ss, aa) counter.update(ss, aa, ns, rr) assert counter.N_sa[ss, aa] == N assert counter.count(ss, aa) == N if rate_power == pytest.approx(1): assert np.allclose(counter.measure(ss, aa), 1.0 / N) elif rate_power == pytest.approx(0.5): assert np.allclose(counter.measure(ss, aa), np.sqrt(1.0 / N)) assert counter.get_n_visited_states() == env.observation_space.n assert np.allclose(counter.get_entropy(), np.log2(env.observation_space.n)) counter.reset()
class UCBVIAgent(IncrementalAgent): """ UCBVI [1]_ with custom exploration bonus. Notes ----- The recommended policy after all the episodes is computed without exploration bonuses. Parameters ---------- env : gym.Env Environment with discrete states and actions. n_episodes : int Number of episodes. gamma : double, default: 1.0 Discount factor in [0, 1]. If gamma is 1.0, the problem is set to be finite-horizon. horizon : int Horizon of the objective function. If None and gamma<1, set to 1/(1-gamma). bonus_scale_factor : double, default: 1.0 Constant by which to multiply the exploration bonus, controls the level of exploration. bonus_type : {"simplified_bernstein"} Type of exploration bonus. Currently, only "simplified_bernstein" is implemented. If `reward_free` is true, this parameter is ignored and the algorithm uses 1/n bonuses. reward_free : bool, default: False If true, ignores rewards and uses only 1/n bonuses. stage_dependent : bool, default: False If true, assume that transitions and rewards can change with the stage h. real_time_dp : bool, default: False If true, uses real-time dynamic programming [2]_ instead of full backward induction for the sampling policy. References ---------- .. [1] Azar et al., 2017 Minimax Regret Bounds for Reinforcement Learning https://arxiv.org/abs/1703.05449 .. [2] Efroni, Yonathan, et al. Tight regret bounds for model-based reinforcement learning with greedy policies. Advances in Neural Information Processing Systems. 2019. https://papers.nips.cc/paper/2019/file/25caef3a545a1fff2ff4055484f0e758-Paper.pdf """ name = "UCBVI" def __init__(self, env, n_episodes=1000, gamma=1.0, horizon=100, bonus_scale_factor=1.0, bonus_type="simplified_bernstein", reward_free=False, stage_dependent=False, real_time_dp=False, **kwargs): # init base class IncrementalAgent.__init__(self, env, **kwargs) self.n_episodes = n_episodes self.gamma = gamma self.horizon = horizon self.bonus_scale_factor = bonus_scale_factor self.bonus_type = bonus_type self.reward_free = reward_free self.stage_dependent = stage_dependent self.real_time_dp = real_time_dp # check environment assert isinstance(self.env.observation_space, spaces.Discrete) assert isinstance(self.env.action_space, spaces.Discrete) # other checks assert gamma >= 0 and gamma <= 1.0 if self.horizon is None: assert gamma < 1.0, \ "If no horizon is given, gamma must be smaller than 1." self.horizon = int(np.ceil(1.0 / (1.0 - gamma))) # maximum value r_range = self.env.reward_range[1] - self.env.reward_range[0] if r_range == np.inf or r_range == 0.0: logger.warning( "{}: Reward range is zero or infinity. ".format(self.name) + "Setting it to 1.") r_range = 1.0 self.v_max = np.zeros(self.horizon) self.v_max[-1] = r_range for hh in reversed(range(self.horizon - 1)): self.v_max[hh] = r_range + self.gamma * self.v_max[hh + 1] # initialize self.reset() def reset(self, **kwargs): H = self.horizon S = self.env.observation_space.n A = self.env.action_space.n if self.stage_dependent: shape_hsa = (H, S, A) shape_hsas = (H, S, A, S) else: shape_hsa = (S, A) shape_hsas = (S, A, S) # (s, a) visit counter self.N_sa = np.zeros(shape_hsa) # (s, a) bonus self.B_sa = np.ones(shape_hsa) # MDP estimator self.R_hat = np.zeros(shape_hsa) self.P_hat = np.ones(shape_hsas) * 1.0 / S # Value functions self.V = np.ones((H, S)) self.Q = np.zeros((H, S, A)) # for rec. policy self.V_policy = np.zeros((H, S)) self.Q_policy = np.zeros((H, S, A)) # Init V and bonus if not self.stage_dependent: self.B_sa *= self.v_max[0] self.V *= self.v_max[0] else: for hh in range(self.horizon): self.B_sa[hh, :, :] = self.v_max[hh] self.V[hh, :] = self.v_max[hh] # ep counter self.episode = 0 # useful object to compute total number of visited states & entropy of visited states self.counter = DiscreteCounter(self.env.observation_space, self.env.action_space) # info self._rewards = np.zeros(self.n_episodes) # update name if self.real_time_dp: self.name = 'UCBVI-RTDP' # default writer self.writer = PeriodicWriter(self.name, log_every=5 * logger.getEffectiveLevel()) def policy(self, state, hh=0, **kwargs): """ Recommended policy. """ assert self.Q_policy is not None return self.Q_policy[hh, state, :].argmax() def _get_action(self, state, hh=0): """ Sampling policy. """ if not self.real_time_dp: assert self.Q is not None return self.Q[hh, state, :].argmax() else: if self.stage_dependent: update_fn = update_value_and_get_action_sd else: update_fn = update_value_and_get_action return update_fn( state, hh, self.V, self.R_hat, self.P_hat, self.B_sa, self.gamma, self.v_max, ) def _compute_bonus(self, n, hh): # reward-free if self.reward_free: bonus = 1.0 / n return bonus # not reward-free if self.bonus_type == "simplified_bernstein": bonus = self.bonus_scale_factor * np.sqrt( 1.0 / n) + self.v_max[hh] / n bonus = min(bonus, self.v_max[hh]) return bonus else: raise ValueError("Error: bonus type {} not implemented".format( self.bonus_type)) def _update(self, state, action, next_state, reward, hh): if self.stage_dependent: self.N_sa[hh, state, action] += 1 nn = self.N_sa[hh, state, action] prev_r = self.R_hat[hh, state, action] prev_p = self.P_hat[hh, state, action, :] self.R_hat[hh, state, action] = (1.0 - 1.0 / nn) * prev_r + reward * 1.0 / nn self.P_hat[hh, state, action, :] = (1.0 - 1.0 / nn) * prev_p self.P_hat[hh, state, action, next_state] += 1.0 / nn self.B_sa[hh, state, action] = self._compute_bonus(nn, hh) else: self.N_sa[state, action] += 1 nn = self.N_sa[state, action] prev_r = self.R_hat[state, action] prev_p = self.P_hat[state, action, :] self.R_hat[state, action] = (1.0 - 1.0 / nn) * prev_r + reward * 1.0 / nn self.P_hat[state, action, :] = (1.0 - 1.0 / nn) * prev_p self.P_hat[state, action, next_state] += 1.0 / nn self.B_sa[state, action] = self._compute_bonus(nn, 0) def _run_episode(self): # interact for H steps episode_rewards = 0 state = self.env.reset() for hh in range(self.horizon): action = self._get_action(state, hh) next_state, reward, done, _ = self.env.step(action) episode_rewards += reward # used for logging only self.counter.update(state, action) if self.reward_free: reward = 0.0 # set to zero before update if reward_free self._update(state, action, next_state, reward, hh) state = next_state if done: break # run backward induction if not self.real_time_dp: if self.stage_dependent: backward_induction_sd(self.Q, self.V, self.R_hat + self.B_sa, self.P_hat, self.gamma, self.v_max[0]) else: backward_induction_in_place(self.Q, self.V, self.R_hat + self.B_sa, self.P_hat, self.horizon, self.gamma, self.v_max[0]) # update info ep = self.episode self._rewards[ep] = episode_rewards self.episode += 1 # writer if self.writer is not None: self.writer.add_scalar("ep reward", episode_rewards, self.episode) self.writer.add_scalar("n_visited_states", self.counter.get_n_visited_states(), self.episode) # return sum of rewards collected in the episode return episode_rewards def partial_fit(self, fraction, **kwargs): assert 0.0 < fraction <= 1.0 n_episodes_to_run = int(np.ceil(fraction * self.n_episodes)) count = 0 while count < n_episodes_to_run and self.episode < self.n_episodes: self._run_episode() count += 1 # compute Q function for the recommended policy if self.stage_dependent: backward_induction_sd(self.Q_policy, self.V_policy, self.R_hat, self.P_hat, self.gamma, self.v_max[0]) else: backward_induction_in_place(self.Q_policy, self.V_policy, self.R_hat, self.P_hat, self.horizon, self.gamma, self.v_max[0]) info = { "n_episodes": self.episode, "episode_rewards": self._rewards[:self.episode] } return info
class OptQLAgent(AgentWithSimplePolicy): """ Optimistic Q-Learning [1]_ with custom exploration bonuses. Parameters ---------- env : gym.Env Environment with discrete states and actions. gamma : double, default: 1.0 Discount factor in [0, 1]. horizon : int Horizon of the objective function. bonus_scale_factor : double, default: 1.0 Constant by which to multiply the exploration bonus, controls the level of exploration. bonus_type : {"simplified_bernstein"} Type of exploration bonus. Currently, only "simplified_bernstein" is implemented. add_bonus_after_update : bool, default: False If True, add bonus to the Q function after performing the update, instead of adding it to the update target. References ---------- .. [1] Jin et al., 2018 Is Q-Learning Provably Efficient? https://arxiv.org/abs/1807.03765 """ name = "OptQL" def __init__(self, env, gamma=1.0, horizon=100, bonus_scale_factor=1.0, bonus_type="simplified_bernstein", add_bonus_after_update=False, **kwargs): # init base class AgentWithSimplePolicy.__init__(self, env, **kwargs) self.gamma = gamma self.horizon = horizon self.bonus_scale_factor = bonus_scale_factor self.bonus_type = bonus_type self.add_bonus_after_update = add_bonus_after_update # check environment assert isinstance(self.env.observation_space, spaces.Discrete) assert isinstance(self.env.action_space, spaces.Discrete) # maximum value r_range = self.env.reward_range[1] - self.env.reward_range[0] if r_range == np.inf or r_range == 0.0: logger.warning( "{}: Reward range is zero or infinity. ".format(self.name) + "Setting it to 1.") r_range = 1.0 self.v_max = np.zeros(self.horizon) self.v_max[-1] = r_range for hh in reversed(range(self.horizon - 1)): self.v_max[hh] = r_range + self.gamma * self.v_max[hh + 1] # initialize self.reset() def reset(self, **kwargs): H = self.horizon S = self.env.observation_space.n A = self.env.action_space.n # (s, a) visit counter self.N_sa = np.zeros((H, S, A)) # Value functions self.V = np.ones((H + 1, S)) self.V[H, :] = 0 self.Q = np.ones((H, S, A)) self.Q_bar = np.ones((H, S, A)) for hh in range(self.horizon): self.V[hh, :] *= self.horizon - hh self.Q[hh, :, :] *= self.horizon - hh self.Q_bar[hh, :, :] *= self.horizon - hh if self.add_bonus_after_update: self.Q *= 0.0 # ep counter self.episode = 0 # useful object to compute total number of visited states & entropy of visited states self.counter = DiscreteCounter(self.env.observation_space, self.env.action_space) def policy(self, observation): """Recommended policy.""" state = observation return self.Q_bar[0, state, :].argmax() def _get_action(self, state, hh=0): """Sampling policy.""" return self.Q_bar[hh, state, :].argmax() def _compute_bonus(self, n, hh): if self.bonus_type == "simplified_bernstein": bonus = self.bonus_scale_factor * np.sqrt( 1.0 / n) + self.v_max[hh] / n bonus = min(bonus, self.v_max[hh]) return bonus else: raise ValueError("Error: bonus type {} not implemented".format( self.bonus_type)) def _update(self, state, action, next_state, reward, hh): self.N_sa[hh, state, action] += 1 nn = self.N_sa[hh, state, action] # learning rate alpha = (self.horizon + 1.0) / (self.horizon + nn) bonus = self._compute_bonus(nn, hh) # bonus in the update if not self.add_bonus_after_update: target = reward + bonus + self.gamma * self.V[hh + 1, next_state] self.Q[hh, state, action] = (1 - alpha) * self.Q[hh, state, action] + alpha * target self.V[hh, state] = min(self.v_max[hh], self.Q[hh, state, :].max()) self.Q_bar[hh, state, action] = self.Q[hh, state, action] # bonus outside the update else: target = reward + self.gamma * self.V[hh + 1, next_state] # bonus not here self.Q[hh, state, action] = (1 - alpha) * self.Q[hh, state, action] + alpha * target self.Q_bar[hh, state, action] = (self.Q[hh, state, action] + bonus ) # bonus here self.V[hh, state] = min(self.v_max[hh], self.Q_bar[hh, state, :].max()) def _run_episode(self): # interact for H steps episode_rewards = 0 state = self.env.reset() for hh in range(self.horizon): action = self._get_action(state, hh) next_state, reward, done, _ = self.env.step(action) episode_rewards += reward # used for logging only self.counter.update(state, action) self._update(state, action, next_state, reward, hh) state = next_state if done: break # update info self.episode += 1 # writer if self.writer is not None: self.writer.add_scalar("episode_rewards", episode_rewards, self.episode) self.writer.add_scalar("n_visited_states", self.counter.get_n_visited_states(), self.episode) # return sum of rewards collected in the episode return episode_rewards def fit(self, budget: int, **kwargs): del kwargs n_episodes_to_run = budget count = 0 while count < n_episodes_to_run: self._run_episode() count += 1
class RLSVIAgent(AgentWithSimplePolicy): """ RLSVI algorithm from [1,2] with Gaussian noise. Notes ----- The recommended policy after all the episodes is computed with the empirical MDP. The std of the noise is of the form: scale/sqrt(n)+ V_max/n as for simplified Bernstein bonuses. Parameters ---------- env : gym.Env Environment with discrete states and actions. gamma : double, default: 1.0 Discount factor in [0, 1]. If gamma is 1.0, the problem is set to be finite-horizon. horizon : int Horizon of the objective function. If None and gamma<1, set to 1/(1-gamma). scale_std_noise : double, delfault: 1.0 scale the std of the noise. At step h the std is scale_std_noise/sqrt(n)+(H-h+1)/n reward_free : bool, default: False If true, ignores rewards. stage_dependent : bool, default: False If true, assume that transitions and rewards can change with the stage h. References ---------- .. [1] Osband et al., 2014 Generalization and Exploration via Randomized Value Functions https://arxiv.org/abs/1402.0635 .. [2] Russo, 2019 Worst-Case Regret Bounds for Exploration via Randomized Value Functions https://arxiv.org/abs/1906.02870 """ name = "RLSVI" def __init__(self, env, gamma=1.0, horizon=100, scale_std_noise=1.0, reward_free=False, stage_dependent=False, **kwargs): # init base class AgentWithSimplePolicy.__init__(self, env, **kwargs) self.gamma = gamma self.horizon = horizon self.scale_std_noise = scale_std_noise self.reward_free = reward_free self.stage_dependent = stage_dependent # check environment assert isinstance(self.env.observation_space, spaces.Discrete) assert isinstance(self.env.action_space, spaces.Discrete) # other checks assert gamma >= 0 and gamma <= 1.0 if self.horizon is None: assert gamma < 1.0, "If no horizon is given, gamma must be smaller than 1." self.horizon = int(np.ceil(1.0 / (1.0 - gamma))) # maximum value r_range = self.env.reward_range[1] - self.env.reward_range[0] if r_range == np.inf or r_range == 0.0: logger.warning( "{}: Reward range is zero or infinity. ".format(self.name) + "Setting it to 1.") r_range = 1.0 self.v_max = np.zeros(self.horizon) self.v_max[-1] = r_range for hh in reversed(range(self.horizon - 1)): self.v_max[hh] = r_range + self.gamma * self.v_max[hh + 1] # initialize self.reset() def reset(self, **kwargs): H = self.horizon S = self.env.observation_space.n A = self.env.action_space.n if self.stage_dependent: shape_hsa = (H, S, A) shape_hsas = (H, S, A, S) else: shape_hsa = (S, A) shape_hsas = (S, A, S) # stds prior self.std1_sa = self.scale_std_noise * np.ones((H, S, A)) self.std2_sa = np.ones((H, S, A)) # visit counter self.N_sa = np.ones(shape_hsa) # MDP estimator self.R_hat = np.zeros(shape_hsa) self.P_hat = np.ones(shape_hsas) * 1.0 / S # Value functions self.V = np.zeros((H, S)) self.Q = np.zeros((H, S, A)) # for rec. policy self.V_policy = np.zeros((H, S)) self.Q_policy = np.zeros((H, S, A)) # Init V and variances for hh in range(self.horizon): self.std2_sa[hh, :, :] *= self.v_max[hh] # ep counter self.episode = 0 # useful object to compute total number of visited states & entropy of visited states self.counter = DiscreteCounter(self.env.observation_space, self.env.action_space) def policy(self, observation): state = observation assert self.Q_policy is not None return self.Q_policy[0, state, :].argmax() def _get_action(self, state, hh=0): """Sampling policy.""" assert self.Q is not None return self.Q[hh, state, :].argmax() def _update(self, state, action, next_state, reward, hh): if self.stage_dependent: self.N_sa[hh, state, action] += 1 nn = self.N_sa[hh, state, action] prev_r = self.R_hat[hh, state, action] prev_p = self.P_hat[hh, state, action, :] self.R_hat[hh, state, action] = (1.0 - 1.0 / nn) * prev_r + reward * 1.0 / nn self.P_hat[hh, state, action, :] = (1.0 - 1.0 / nn) * prev_p self.P_hat[hh, state, action, next_state] += 1.0 / nn else: self.N_sa[state, action] += 1 nn = self.N_sa[state, action] prev_r = self.R_hat[state, action] prev_p = self.P_hat[state, action, :] self.R_hat[state, action] = (1.0 - 1.0 / nn) * prev_r + reward * 1.0 / nn self.P_hat[state, action, :] = (1.0 - 1.0 / nn) * prev_p self.P_hat[state, action, next_state] += 1.0 / nn def _run_episode(self): # interact for H steps episode_rewards = 0 # stds scale/sqrt(n)+(H-h+1)/n std_sa = self.std1_sa / np.sqrt(self.N_sa) + self.std2_sa / self.N_sa noise_sa = self.rng.normal(self.R_hat, std_sa) # run backward noisy induction if self.stage_dependent: backward_induction_sd( self.Q, self.V, self.R_hat + noise_sa, self.P_hat, self.gamma, self.v_max[0], ) else: backward_induction_reward_sd( self.Q, self.V, self.R_hat + noise_sa, self.P_hat, self.gamma, self.v_max[0], ) state = self.env.reset() for hh in range(self.horizon): action = self._get_action(state, hh) next_state, reward, done, _ = self.env.step(action) episode_rewards += reward # used for logging only self.counter.update(state, action) if self.reward_free: reward = 0.0 # set to zero before update if reward_free self._update(state, action, next_state, reward, hh) state = next_state if done: break # update info self.episode += 1 # writer if self.writer is not None: self.writer.add_scalar("episode_rewards", episode_rewards, self.episode) self.writer.add_scalar("n_visited_states", self.counter.get_n_visited_states(), self.episode) # return sum of rewards collected in the episode return episode_rewards def fit(self, budget: int, **kwargs): del kwargs n_episodes_to_run = budget count = 0 while count < n_episodes_to_run: self._run_episode() count += 1 # compute Q function for the recommended policy if self.stage_dependent: backward_induction_sd( self.Q_policy, self.V_policy, self.R_hat, self.P_hat, self.gamma, self.v_max[0], ) else: backward_induction_in_place( self.Q_policy, self.V_policy, self.R_hat, self.P_hat, self.horizon, self.gamma, self.v_max[0], )
class PSRLAgent(AgentWithSimplePolicy): """ PSRL algorithm from [1] with beta prior for the "Bernoullized" rewards (instead of Gaussian-gamma prior). Notes ----- The recommended policy after all the episodes is computed without exploration bonuses. Parameters ---------- env : gym.Env Environment with discrete states and actions. gamma : double, default: 1.0 Discount factor in [0, 1]. If gamma is 1.0, the problem is set to be finite-horizon. horizon : int Horizon of the objective function. If None and gamma<1, set to 1/(1-gamma). scale_prior_reward : double, delfault: 1.0 scale of the Beta (uniform) prior, i.e prior is Beta(scale_prior_reward*(1,1)) scale_prior_transition : double, default: 1/number of state scale of the (uniform) Dirichlet prior, i.e prior is Dirichlet(scale_prior_transition*(1,...,1)) bernoullized_reward: bool, default: True If true the rewards are Bernoullized reward_free : bool, default: False If true, ignores rewards and uses only 1/n bonuses. stage_dependent : bool, default: False If true, assume that transitions and rewards can change with the stage h. References ---------- .. [1] Osband et al., 2013 (More) Efficient Reinforcement Learning via Posterior Sampling https://arxiv.org/abs/1306.0940 """ name = "PSRL" def __init__(self, env, gamma=1.0, horizon=100, scale_prior_reward=1, scale_prior_transition=None, bernoullized_reward=True, reward_free=False, stage_dependent=False, **kwargs): # init base class AgentWithSimplePolicy.__init__(self, env, **kwargs) self.gamma = gamma self.horizon = horizon self.scale_prior_reward = scale_prior_reward self.scale_prior_transition = scale_prior_transition if scale_prior_transition is None: self.scale_prior_transition = 1.0 / self.env.observation_space.n self.bernoullized_reward = bernoullized_reward self.reward_free = reward_free self.stage_dependent = stage_dependent # check environment assert isinstance(self.env.observation_space, spaces.Discrete) assert isinstance(self.env.action_space, spaces.Discrete) # other checks assert gamma >= 0 and gamma <= 1.0 if self.horizon is None: assert gamma < 1.0, "If no horizon is given, gamma must be smaller than 1." self.horizon = int(np.ceil(1.0 / (1.0 - gamma))) # maximum value r_range = self.env.reward_range[1] - self.env.reward_range[0] if r_range == np.inf or r_range == 0.0: logger.warning( "{}: Reward range is zero or infinity. ".format(self.name) + "Setting it to 1.") r_range = 1.0 self.v_max = np.zeros(self.horizon) self.v_max[-1] = r_range for hh in reversed(range(self.horizon - 1)): self.v_max[hh] = r_range + self.gamma * self.v_max[hh + 1] # initialize self.reset() def reset(self, **kwargs): H = self.horizon S = self.env.observation_space.n A = self.env.action_space.n if self.stage_dependent: shape_hsa = (H, S, A) shape_hsas = (H, S, A, S) else: shape_hsa = (S, A) shape_hsas = (S, A, S) # Prior transitions self.N_sas = self.scale_prior_transition * np.ones(shape_hsas) # Prior rewards self.M_sa = self.scale_prior_reward * np.ones(shape_hsa + (2, )) # Value functions self.V = np.zeros((H, S)) self.Q = np.zeros((H, S, A)) # for rec. policy self.V_policy = np.zeros((H, S)) self.Q_policy = np.zeros((H, S, A)) # ep counter self.episode = 0 # useful object to compute total number of visited states & entropy of visited states self.counter = DiscreteCounter(self.env.observation_space, self.env.action_space) def policy(self, observation): state = observation assert self.Q_policy is not None return self.Q_policy[0, state, :].argmax() def _get_action(self, state, hh=0): """Sampling policy.""" assert self.Q is not None return self.Q[hh, state, :].argmax() def _update(self, state, action, next_state, reward, hh): bern_reward = reward if self.bernoullized_reward: bern_reward = self.rng.binomial(1, reward) # update posterior if self.stage_dependent: self.N_sas[hh, state, action, next_state] += 1 self.M_sa[hh, state, action, 0] += bern_reward self.M_sa[hh, state, action, 1] += 1 - bern_reward else: self.N_sas[state, action, next_state] += 1 self.M_sa[state, action, 0] += bern_reward self.M_sa[state, action, 1] += 1 - bern_reward def _run_episode(self): # sample reward and transitions from posterior self.R_sample = self.rng.beta(self.M_sa[..., 0], self.M_sa[..., 1]) self.P_sample = self.rng.gamma(self.N_sas) self.P_sample = self.P_sample / self.P_sample.sum(-1, keepdims=True) # run backward induction if self.stage_dependent: backward_induction_sd(self.Q, self.V, self.R_sample, self.P_sample, self.gamma, self.v_max[0]) else: backward_induction_in_place( self.Q, self.V, self.R_sample, self.P_sample, self.horizon, self.gamma, self.v_max[0], ) # interact for H steps episode_rewards = 0 state = self.env.reset() for hh in range(self.horizon): action = self._get_action(state, hh) next_state, reward, done, _ = self.env.step(action) episode_rewards += reward # used for logging only self.counter.update(state, action) if self.reward_free: reward = 0.0 # set to zero before update if reward_free self._update(state, action, next_state, reward, hh) state = next_state if done: break # update info self.episode += 1 # writer if self.writer is not None: self.writer.add_scalar("episode_rewards", episode_rewards, self.episode) self.writer.add_scalar("n_visited_states", self.counter.get_n_visited_states(), self.episode) # return sum of rewards collected in the episode return episode_rewards def fit(self, budget: int, **kwargs): del kwargs n_episodes_to_run = budget count = 0 while count < n_episodes_to_run: self._run_episode() count += 1 # compute Q function for the recommended policy R_hat = self.M_sa[..., 0] / (self.M_sa[..., 0] + self.M_sa[..., 1]) P_hat = self.N_sas / self.N_sas.sum(-1, keepdims=True) if self.stage_dependent: backward_induction_sd(self.Q_policy, self.V_policy, R_hat, P_hat, self.gamma, self.v_max[0]) else: backward_induction_in_place( self.Q_policy, self.V_policy, R_hat, P_hat, self.horizon, self.gamma, self.v_max[0], )