def test_simple_q_loss_function(self): """Tests the Simple-Q loss function results on all frameworks.""" config = dqn.simple_q.SimpleQConfig().rollouts(num_rollout_workers=0) # Use very simple net (layer0=10 nodes, q-layer=2 nodes (2 actions)). config.training(model={ "fcnet_hiddens": [10], "fcnet_activation": "linear", }) for fw in framework_iterator(config): # Generate Trainer and get its default Policy object. trainer = dqn.SimpleQTrainer(config=config, env="CartPole-v0") policy = trainer.get_policy() # Batch of size=2. input_ = SampleBatch({ SampleBatch.CUR_OBS: np.random.random(size=(2, 4)), SampleBatch.ACTIONS: np.array([0, 1]), SampleBatch.REWARDS: np.array([0.4, -1.23]), SampleBatch.DONES: np.array([False, False]), SampleBatch.NEXT_OBS: np.random.random(size=(2, 4)), SampleBatch.EPS_ID: np.array([1234, 1234]), SampleBatch.AGENT_INDEX: np.array([0, 0]), SampleBatch.ACTION_LOGP: np.array([-0.1, -0.1]), SampleBatch.ACTION_DIST_INPUTS: np.array([[0.1, 0.2], [-0.1, -0.2]]), SampleBatch.ACTION_PROB: np.array([0.1, 0.2]), "q_values": np.array([[0.1, 0.2], [0.2, 0.1]]), }) # Get model vars for computing expected model outs (q-vals). # 0=layer-kernel; 1=layer-bias; 2=q-val-kernel; 3=q-val-bias vars = policy.get_weights() if isinstance(vars, dict): vars = list(vars.values()) vars_t = policy.target_model.variables() if fw == "tf": vars_t = policy.get_session().run(vars_t) # Q(s,a) outputs. q_t = np.sum( one_hot(input_[SampleBatch.ACTIONS], 2) * fc( fc( input_[SampleBatch.CUR_OBS], vars[0 if fw != "torch" else 2], vars[1 if fw != "torch" else 3], framework=fw, ), vars[2 if fw != "torch" else 0], vars[3 if fw != "torch" else 1], framework=fw, ), 1, ) # max[a'](Qtarget(s',a')) outputs. q_target_tp1 = np.max( fc( fc( input_[SampleBatch.NEXT_OBS], vars_t[0 if fw != "torch" else 2], vars_t[1 if fw != "torch" else 3], framework=fw, ), vars_t[2 if fw != "torch" else 0], vars_t[3 if fw != "torch" else 1], framework=fw, ), 1, ) # TD-errors (Bellman equation). td_error = q_t - config.gamma * input_[ SampleBatch.REWARDS] + q_target_tp1 # Huber/Square loss on TD-error. expected_loss = huber_loss(td_error).mean() if fw == "torch": input_ = policy._lazy_tensor_dict(input_) # Get actual out and compare. if fw == "tf": out = policy.get_session().run( policy._loss, feed_dict=policy._get_loss_inputs_dict(input_, shuffle=False), ) else: out = (loss_torch if fw == "torch" else loss_tf)(policy, policy.model, None, input_) check(out, expected_loss, decimals=1)
def _ddpg_loss_helper(self, train_batch, weights, ks, fw, gamma, huber_threshold, l2_reg, sess): """Emulates DDPG loss functions for tf and torch.""" model_out_t = train_batch[SampleBatch.CUR_OBS] target_model_out_tp1 = train_batch[SampleBatch.NEXT_OBS] # get_policy_output policy_t = sigmoid(2.0 * fc( relu(fc(model_out_t, weights[ks[1]], weights[ks[0]], framework=fw)), weights[ks[5]], weights[ks[4]])) # Get policy output for t+1 (target model). policy_tp1 = sigmoid(2.0 * fc( relu( fc(target_model_out_tp1, weights[ks[3]], weights[ks[2]], framework=fw)), weights[ks[7]], weights[ks[6]])) # Assume no smooth target policy. policy_tp1_smoothed = policy_tp1 # Q-values for the actually selected actions. # get_q_values q_t = fc(relu( fc(np.concatenate([model_out_t, train_batch[SampleBatch.ACTIONS]], -1), weights[ks[9]], weights[ks[8]], framework=fw)), weights[ks[11]], weights[ks[10]], framework=fw) twin_q_t = fc(relu( fc(np.concatenate([model_out_t, train_batch[SampleBatch.ACTIONS]], -1), weights[ks[13]], weights[ks[12]], framework=fw)), weights[ks[15]], weights[ks[14]], framework=fw) # Q-values for current policy in given current state. # get_q_values q_t_det_policy = fc(relu( fc(np.concatenate([model_out_t, policy_t], -1), weights[ks[9]], weights[ks[8]], framework=fw)), weights[ks[11]], weights[ks[10]], framework=fw) # Target q network evaluation. # target_model.get_q_values q_tp1 = fc(relu( fc(np.concatenate([target_model_out_tp1, policy_tp1_smoothed], -1), weights[ks[17]], weights[ks[16]], framework=fw)), weights[ks[19]], weights[ks[18]], framework=fw) twin_q_tp1 = fc(relu( fc(np.concatenate([target_model_out_tp1, policy_tp1_smoothed], -1), weights[ks[21]], weights[ks[20]], framework=fw)), weights[ks[23]], weights[ks[22]], framework=fw) q_t_selected = np.squeeze(q_t, axis=-1) twin_q_t_selected = np.squeeze(twin_q_t, axis=-1) q_tp1 = np.minimum(q_tp1, twin_q_tp1) q_tp1_best = np.squeeze(q_tp1, axis=-1) dones = train_batch[SampleBatch.DONES] rewards = train_batch[SampleBatch.REWARDS] if fw == "torch": dones = dones.float().numpy() rewards = rewards.numpy() q_tp1_best_masked = (1.0 - dones) * q_tp1_best q_t_selected_target = rewards + gamma * q_tp1_best_masked td_error = q_t_selected - q_t_selected_target twin_td_error = twin_q_t_selected - q_t_selected_target td_error = td_error + twin_td_error errors = huber_loss(td_error, huber_threshold) + \ huber_loss(twin_td_error, huber_threshold) critic_loss = np.mean(errors) actor_loss = -np.mean(q_t_det_policy) # Add l2-regularization if required. for name, var in weights.items(): if re.match("default_policy/actor_(hidden_0|out)/kernel", name): actor_loss += (l2_reg * l2_loss(var)) elif re.match("default_policy/sequential(_1)?/\\w+/kernel", name): critic_loss += (l2_reg * l2_loss(var)) return critic_loss, actor_loss, td_error
def _sac_loss_helper(self, train_batch, weights, ks, log_alpha, fw, gamma, sess): """Emulates SAC loss functions for tf and torch.""" # ks: # 0=log_alpha # 1=target log-alpha (not used) # 2=action hidden bias # 3=action hidden kernel # 4=action out bias # 5=action out kernel # 6=Q hidden bias # 7=Q hidden kernel # 8=Q out bias # 9=Q out kernel # 14=target Q hidden bias # 15=target Q hidden kernel # 16=target Q out bias # 17=target Q out kernel alpha = np.exp(log_alpha) # cls = TorchSquashedGaussian if fw == "torch" else SquashedGaussian cls = TorchDirichlet if fw == "torch" else Dirichlet model_out_t = train_batch[SampleBatch.CUR_OBS] model_out_tp1 = train_batch[SampleBatch.NEXT_OBS] target_model_out_tp1 = train_batch[SampleBatch.NEXT_OBS] # get_policy_output action_dist_t = cls( fc( relu( fc(model_out_t, weights[ks[1]], weights[ks[0]], framework=fw)), weights[ks[9]], weights[ks[8]]), None) policy_t = action_dist_t.deterministic_sample() log_pis_t = action_dist_t.logp(policy_t) if sess: log_pis_t = sess.run(log_pis_t) policy_t = sess.run(policy_t) log_pis_t = np.expand_dims(log_pis_t, -1) # Get policy output for t+1. action_dist_tp1 = cls( fc( relu( fc(model_out_tp1, weights[ks[1]], weights[ks[0]], framework=fw)), weights[ks[9]], weights[ks[8]]), None) policy_tp1 = action_dist_tp1.deterministic_sample() log_pis_tp1 = action_dist_tp1.logp(policy_tp1) if sess: log_pis_tp1 = sess.run(log_pis_tp1) policy_tp1 = sess.run(policy_tp1) log_pis_tp1 = np.expand_dims(log_pis_tp1, -1) # Q-values for the actually selected actions. # get_q_values q_t = fc(relu( fc(np.concatenate([model_out_t, train_batch[SampleBatch.ACTIONS]], -1), weights[ks[3]], weights[ks[2]], framework=fw)), weights[ks[11]], weights[ks[10]], framework=fw) # Q-values for current policy in given current state. # get_q_values q_t_det_policy = fc(relu( fc(np.concatenate([model_out_t, policy_t], -1), weights[ks[3]], weights[ks[2]], framework=fw)), weights[ks[11]], weights[ks[10]], framework=fw) # Target q network evaluation. # target_model.get_q_values if fw == "tf": q_tp1 = fc(relu( fc(np.concatenate([target_model_out_tp1, policy_tp1], -1), weights[ks[7]], weights[ks[6]], framework=fw)), weights[ks[15]], weights[ks[14]], framework=fw) else: assert fw == "tfe" q_tp1 = fc(relu( fc(np.concatenate([target_model_out_tp1, policy_tp1], -1), weights[ks[7]], weights[ks[6]], framework=fw)), weights[ks[9]], weights[ks[8]], framework=fw) q_t_selected = np.squeeze(q_t, axis=-1) q_tp1 -= alpha * log_pis_tp1 q_tp1_best = np.squeeze(q_tp1, axis=-1) dones = train_batch[SampleBatch.DONES] rewards = train_batch[SampleBatch.REWARDS] if fw == "torch": dones = dones.float().numpy() rewards = rewards.numpy() q_tp1_best_masked = (1.0 - dones) * q_tp1_best q_t_selected_target = rewards + gamma * q_tp1_best_masked base_td_error = np.abs(q_t_selected - q_t_selected_target) td_error = base_td_error critic_loss = [ np.mean(train_batch["weights"] * huber_loss(q_t_selected_target - q_t_selected)) ] target_entropy = -np.prod((1, )) alpha_loss = -np.mean(log_alpha * (log_pis_t + target_entropy)) actor_loss = np.mean(alpha * log_pis_t - q_t_det_policy) return critic_loss, actor_loss, alpha_loss, td_error
def test_simple_q_loss_function(self): """Tests the Simple-Q loss function results on all frameworks.""" config = dqn.SIMPLE_Q_DEFAULT_CONFIG.copy() # Run locally. config["num_workers"] = 0 # Use very simple net (layer0=10 nodes, q-layer=2 nodes (2 actions)). config["model"]["fcnet_hiddens"] = [10] config["model"]["fcnet_activation"] = "linear" for fw in framework_iterator(config): # Generate Trainer and get its default Policy object. trainer = dqn.SimpleQTrainer(config=config, env="CartPole-v0") policy = trainer.get_policy() # Batch of size=2. input_ = { SampleBatch.CUR_OBS: np.random.random(size=(2, 4)), SampleBatch.ACTIONS: np.array([0, 1]), SampleBatch.REWARDS: np.array([0.4, -1.23]), SampleBatch.DONES: np.array([False, False]), SampleBatch.NEXT_OBS: np.random.random(size=(2, 4)) } # Get model vars for computing expected model outs (q-vals). # 0=layer-kernel; 1=layer-bias; 2=q-val-kernel; 3=q-val-bias vars = policy.get_weights() if isinstance(vars, dict): vars = list(vars.values()) vars_t = policy.target_q_func_vars if fw == "tf": vars_t = policy.get_session().run(vars_t) # Q(s,a) outputs. q_t = np.sum( one_hot(input_[SampleBatch.ACTIONS], 2) * fc(fc(input_[SampleBatch.CUR_OBS], vars[0 if fw != "torch" else 2], vars[1 if fw != "torch" else 3], framework=fw), vars[2 if fw != "torch" else 0], vars[3 if fw != "torch" else 1], framework=fw), 1) # max[a'](Qtarget(s',a')) outputs. q_target_tp1 = np.max( fc(fc(input_[SampleBatch.NEXT_OBS], vars_t[0 if fw != "torch" else 2], vars_t[1 if fw != "torch" else 3], framework=fw), vars_t[2 if fw != "torch" else 0], vars_t[3 if fw != "torch" else 1], framework=fw), 1) # TD-errors (Bellman equation). td_error = q_t - config["gamma"] * input_[SampleBatch.REWARDS] + \ q_target_tp1 # Huber/Square loss on TD-error. expected_loss = huber_loss(td_error).mean() if fw == "torch": input_ = policy._lazy_tensor_dict(input_) # Get actual out and compare. if fw == "tf": out = policy.get_session().run( policy._loss, feed_dict=policy._get_loss_inputs_dict(input_, shuffle=False)) else: out = (loss_torch if fw == "torch" else loss_tf)(policy, policy.model, None, input_) check(out, expected_loss, decimals=1)