def process_batch(engine_for_batch, batch): optimizer.zero_grad() loss_v = utils.calc_loss_dqn(batch, net, target_net.target_model, gamma=params.gamma, device=device) loss_v.backward() optimizer.step() epsilon_tracker.frame(engine_for_batch.state.iteration) if engine_for_batch.state.iteration % params.target_net_sync == 0: target_net.sync() if engine.state.iteration % EVAL_EVERY_FRAME == 0: eval_states = getattr(engine.state, "eval_states", None) if eval_states is None: eval_states = buffer.sample(STATES_TO_EVALUATE) eval_states = [np.array(transition.state, copy=False) for transition in eval_states] eval_states = np.array(eval_states, copy=False) engine.state.eval_states = eval_states evaluate_states(eval_states, net, device, engine) return { "loss": loss_v.item(), "epsilon": selector.epsilon, }
def process_batch(engine_for_batch, batch_data): batch, batch_indices, batch_weights = batch_data optimizer.zero_grad() loss_v, sample_priority = calc_loss(batch, batch_weights, net, target_net.target_model, gamma=params.gamma, _device=str(device)) loss_v.backward() optimizer.step() buffer.update_priorities(batch_indices, sample_priority) epsilon_tracker.frame(engine_for_batch.state.iteration) if engine_for_batch.state.iteration % params.target_net_sync == 0: target_net.sync() return { "loss": loss_v.item(), "epsilon": selector.epsilon, "beta": buffer.update_beta(engine.state.iteration), }
def process_batch(engine_for_batch, batch): optimizer.zero_grad() loss_v = utils.calc_loss_dqn(batch, net, target_net.target_model, gamma=params.gamma, device=device) loss_v.backward() optimizer.step() epsilon_tracker.frame(engine_for_batch.state.iteration) if engine_for_batch.state.iteration % params.target_net_sync == 0: target_net.sync() return { "loss": loss_v.item(), "epsilon": selector.epsilon, }