Method train

examples/reinforcement_learning/tutorial_C51.py:228–253

(self, b_o, b_a, b_r, b_o_, b_d)

        return self.qnet(obv)

    def train(self, b_o, b_a, b_r, b_o_, b_d):
        # TODO: move q_estimation in tf.function
        b_dist_ = np.exp(self.targetqnet(b_o_).numpy())
        b_a_ = (b_dist_ * vrange).sum(-1).argmax(1)
        b_tzj = np.clip(reward_gamma * (1 - b_d[:, None]) * vrange[None, :] + b_r[:, None], min_value, max_value)
        b_i = (b_tzj - min_value) / deltaz
        b_l = np.floor(b_i).astype('int64')
        b_u = np.ceil(b_i).astype('int64')
        templ = b_dist_[range(batch_size), b_a_, :] * (b_u - b_i)
        tempu = b_dist_[range(batch_size), b_a_, :] * (b_i - b_l)
        b_m = np.zeros((batch_size, atom_num))
        # TODO: aggregate value by index and batch update (scatter_add)
        for j in range(batch_size):
            for k in range(atom_num):
                b_m[j][b_l[j][k]] += templ[j][k]
                b_m[j][b_u[j][k]] += tempu[j][k]
        b_m = tf.convert_to_tensor(b_m, dtype='float32')
        b_index = np.stack([range(batch_size), b_a], 1)
        b_index = tf.convert_to_tensor(b_index, 'int64')

        self._train_func(b_o, b_index, b_m)

        self.niter += 1
        if self.niter % target_q_update_freq == 0:
            sync(self.qnet, self.targetqnet)
            self.save(args.save_path)

    def save(self, path):
        if path is None: