Пример #1
0
def exploration(args, actor_id, param_queue):
    writer = SummaryWriter(comment="-{}-eval".format(args.env))

    args.clip_rewards = False
    args.episode_life = False
    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + actor_id
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env, args)

    param = param_queue.get(block=True)
    model.load_state_dict(param)
    param = None
    print("Received First Parameter!")

    episode_reward, episode_length, episode_idx = 0, 0, 0
    state = env.reset()
    tb_dict = {k: [] for k in ['episode_reward', 'episode_length']}
    while True:
        action, _ = model.act(torch.FloatTensor(np.array(state)), 0.)
        next_state, reward, done, _ = env.step(action)

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done or episode_length == args.max_episode_length:
            state = env.reset()
            tb_dict["episode_reward"].append(episode_reward)
            tb_dict["episode_length"].append(episode_length)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1
            param = param_queue.get()
            model.load_state_dict(param)
            print(f"{datetime.now()} Updated Parameter..")

            if (episode_idx *
                    args.num_envs_per_worker) % args.tb_interval == 0:
                writer.add_scalar('evaluator/episode_reward_mean',
                                  np.mean(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/episode_reward_max',
                                  np.max(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/episode_reward_min',
                                  np.min(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/episode_reward_std',
                                  np.std(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/episode_length_mean',
                                  np.mean(tb_dict['episode_length']),
                                  episode_idx)
                tb_dict['episode_reward'].clear()
                tb_dict['episode_length'].clear()
Пример #2
0
def train(args, n_actors, batch_queue, prios_queue, param_queue):
    env = wrapper.make_atari(args.env)
    env = wrapper.wrap_atari_dqn(env, args)
    utils.set_global_seeds(args.seed, use_torch=True)

    model = DuelingDQN(env, args).to(args.device)
    # model.load_state_dict(torch.load('model_30h.pth'))
    tgt_model = DuelingDQN(env, args).to(args.device)
    tgt_model.load_state_dict(model.state_dict())

    writer = SummaryWriter(comment="-{}-learner".format(args.env))
    optimizer = torch.optim.Adam(model.parameters(), args.lr)
    # optimizer = torch.optim.RMSprop(model.parameters(), args.lr, alpha=0.95, eps=1.5e-7, centered=True)

    check_connection(n_actors)

    param_queue.put(model.state_dict())
    learn_idx = 0
    ts = time.time()
    tb_dict = {
        k: []
        for k in ['loss', 'grad_norm', 'max_q', 'mean_q', 'min_q']
    }
    while True:
        *batch, idxes = batch_queue.get()
        loss, prios, q_values = utils.compute_loss(model, tgt_model, batch,
                                                   args.n_steps, args.gamma)
        grad_norm = utils.update_parameters(loss, model, optimizer,
                                            args.max_norm)
        prios_queue.put((idxes, prios))
        batch, idxes, prios = None, None, None
        learn_idx += 1

        tb_dict["loss"].append(float(loss))
        tb_dict["grad_norm"].append(float(grad_norm))
        tb_dict["max_q"].append(float(torch.max(q_values)))
        tb_dict["mean_q"].append(float(torch.mean(q_values)))
        tb_dict["min_q"].append(float(torch.min(q_values)))

        if args.soft_target_update:
            tau = args.tau
            for p_tgt, p in zip(tgt_model.parameters(), model.parameters()):
                p_tgt.data *= 1 - tau
                p_tgt.data += tau * p
        elif learn_idx % args.target_update_interval == 0:
            print("Updating Target Network..")
            tgt_model.load_state_dict(model.state_dict())
        if learn_idx % args.save_interval == 0:
            print("Saving Model..")
            torch.save(model.state_dict(), "model.pth")
        if learn_idx % args.publish_param_interval == 0:
            param_queue.put(model.state_dict())
        if learn_idx % args.tb_interval == 0:
            bps = args.tb_interval / (time.time() - ts)
            print("Step: {:8} / BPS: {:.2f}".format(learn_idx, bps))
            writer.add_scalar("learner/BPS", bps, learn_idx)
            for k, v in tb_dict.items():
                writer.add_scalar(f'learner/{k}', np.mean(v), learn_idx)
                v.clear()
            ts = time.time()
Пример #3
0
def main():
    args = argparser()

    args.clip_rewards = False
    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + 1122
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env)
    model.load_state_dict(torch.load('model.pth', map_location='cpu'))

    episode_reward, episode_length = 0, 0
    state = env.reset()
    while True:
        if args.render:
            env.render()
        action, _ = model.act(torch.FloatTensor(np.array(state)), 0.)
        next_state, reward, done, _ = env.step(action)

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done:
            state = env.reset()
            print("Episode Length / Reward: {} / {}".format(
                episode_length, episode_reward))
            episode_reward = 0
            episode_length = 0
Пример #4
0
def exploration(args, actor_id, n_actors, param_queue, send_queue,
                req_param_queue):
    writer = SummaryWriter(comment="-{}-actor{}".format(args.env, actor_id))

    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + actor_id
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env)
    epsilon = args.eps_base**(1 + actor_id / (n_actors - 1) * args.eps_alpha)
    storage = BatchStorage(args.n_steps, args.gamma)
    req_param_queue.put(True)
    param = param_queue.get(block=True)
    model.load_state_dict(param)
    param = None
    print("Received First Parameter!")

    episode_reward, episode_length, episode_idx, actor_idx = 0, 0, 0, 0
    state = env.reset()
    while True:
        action, q_values = model.act(torch.FloatTensor(np.array(state)),
                                     epsilon)
        next_state, reward, done, _ = env.step(action)
        com_state = zlib.compress(np.array(state).tobytes())
        storage.add(com_state, reward, action, done, q_values)

        state = next_state
        episode_reward += reward
        episode_length += 1
        actor_idx += 1

        if done or episode_length == args.max_episode_length:
            state = env.reset()
            writer.add_scalar("actor/episode_reward", episode_reward,
                              episode_idx)
            writer.add_scalar("actor/episode_length", episode_length,
                              episode_idx)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1

        if actor_idx % args.update_interval == 0:
            try:
                req_param_queue.put(True)
                param = param_queue.get(block=True)
                model.load_state_dict(param)
                print("Updated Parameter..")
            except queue.Empty:
                pass

        if len(storage) == args.send_interval:
            batch, prios = storage.make_batch()
            send_queue.put((batch, prios))
            batch, prios = None, None
            storage.reset()
Пример #5
0
def train(args, n_actors, batch_queue, prios_queue, param_queue):
    env = RunTagEnv(width=5,
                    height=5,
                    number_of_subordinates=1,
                    max_steps=1000)
    #env = wrapper.make_atari(args.env)
    #env = wrapper.wrap_atari_dqn(env, args)
    utils.set_global_seeds(args.seed, use_torch=True)

    model = DuelingDQN(env).to(args.device)
    tgt_model = DuelingDQN(env).to(args.device)
    tgt_model.load_state_dict(model.state_dict())

    writer = SummaryWriter(comment="-{}-learner".format(args.env))
    # optimizer = torch.optim.Adam(model.parameters(), args.lr)
    optimizer = torch.optim.RMSprop(model.parameters(),
                                    args.lr,
                                    alpha=0.95,
                                    eps=1.5e-7,
                                    centered=True)

    check_connection(n_actors)

    param_queue.put(model.state_dict())
    learn_idx = 0
    ts = time.time()
    while True:
        *batch, idxes = batch_queue.get()
        loss, prios = utils.compute_loss(model, tgt_model, batch, args.n_steps,
                                         args.gamma)
        grad_norm = utils.update_parameters(loss, model, optimizer,
                                            args.max_norm)
        print('Updated parameters!')
        prios_queue.put((idxes, prios))
        batch, idxes, prios = None, None, None
        learn_idx += 1

        writer.add_scalar("learner/loss", loss, learn_idx)
        writer.add_scalar("learner/grad_norm", grad_norm, learn_idx)

        if learn_idx % args.target_update_interval == 0:
            print("Updating Target Network..")
            tgt_model.load_state_dict(model.state_dict())
        if learn_idx % args.save_interval == 0:
            print("Saving Model..")
            torch.save(model.state_dict(), "model.pth")
        if learn_idx % args.publish_param_interval == 0:
            param_queue.put(model.state_dict())
        if learn_idx % args.bps_interval == 0:
            bps = args.bps_interval / (time.time() - ts)
            print("Step: {:8} / BPS: {:.2f}".format(learn_idx, bps))
            writer.add_scalar("learner/BPS", bps, learn_idx)
            ts = time.time()
Пример #6
0
def main():
    learner_ip = get_environ()
    args = argparser()

    writer = SummaryWriter(comment="-{}-eval".format(args.env))

    ctx = zmq.Context()
    param_socket = ctx.socket(zmq.SUB)
    param_socket.setsockopt(zmq.SUBSCRIBE, b'')
    param_socket.setsockopt(zmq.CONFLATE, 1)
    param_socket.connect('tcp://{}:52001'.format(learner_ip))

    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + 1122
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env)

    data = param_socket.recv(copy=False)
    param = pickle.loads(data)
    model.load_state_dict(param)
    print("Loaded first parameter from learner")

    episode_reward, episode_length, episode_idx = 0, 0, 0
    state = env.reset()
    while True:
        if args.render:
            env.render()
        action, _ = model.act(torch.FloatTensor(np.array(state)), 0.01)
        next_state, reward, done, _ = env.step(action)

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done:
            state = env.reset()
            writer.add_scalar("eval/episode_reward", episode_reward,
                              episode_idx)
            writer.add_scalar("eval/episode_length", episode_length,
                              episode_idx)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1

            if episode_idx % args.eval_update_interval == 0:
                data = param_socket.recv(copy=False)
                param = pickle.loads(data)
                model.load_state_dict(param)
Пример #7
0
def main():
    args = argparser()

    args.clip_rewards = False
    args.episode_life=False
    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    # seed = args.seed + 1122
    # utils.set_global_seeds(seed, use_torch=True)
    # env.seed(seed)

    model = DuelingDQN(env, args)
    model.load_state_dict(torch.load('model.pth', map_location='cpu'))

    episode_reward, episode_length = 0, 0
    state = env.reset()
    if not os.path.exists('plays'):
        os.mkdir('plays')
    video = cv2.VideoWriter('plays/tmp.avi', cv2.VideoWriter_fourcc(*'DIVX'), 15, (160, 210))
    while True:
        img = env.render(mode='rgb_array')
        model.zero_grad()
        state = torch.tensor(state[np.newaxis, :], dtype=torch.float32, requires_grad=True)
        value, action = model(state).max(1)
        value = value[0]
        action = action[0]
        value.backward()
        img_gradient = np.abs(state.grad.numpy())
        img_gradient = np.sum(img_gradient, axis=(0,1))
        img_gradient = (img_gradient - np.min(img_gradient)) / (np.max(img_gradient) - np.min(img_gradient))
        img_gradient = img_gradient.transpose()
        img_gradient = cv2.resize(img_gradient, (160, 210))[...,np.newaxis]
        img_gradient = img_gradient * 255
        masked_img = (img + img_gradient).astype(np.uint8)
        masked_img = np.clip(masked_img, 0, 255)
        video.write(masked_img)
        next_state, reward, done, _ = env.step(int(action))

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done:
            state = env.reset()
            print("Episode Length / Reward: {} / {}".format(episode_length, episode_reward))
            video.release()
            os.rename('plays/tmp.avi', f'plays/{args.env}-{episode_reward}.avi')
            video = cv2.VideoWriter('plays/tmp.avi', cv2.VideoWriter_fourcc(*'DIVX'), 15, (160, 210))
            episode_reward = 0
            episode_length = 0
Пример #8
0
def exploration_eval(args, actor_id, param_queue):
    writer = SummaryWriter(comment="-{}-eval".format(args.env))

    args.clip_rewards = False
    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + actor_id
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env)

    param = param_queue.get(block=True)
    model.load_state_dict(param)
    param = None
    print("Received First Parameter!")

    episode_reward, episode_length, episode_idx = 0, 0, 0
    state = env.reset()
    while True:
        action, _ = model.act(torch.FloatTensor(np.array(state)), 0.)
        next_state, reward, done, _ = env.step(action)

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done or episode_length == args.max_episode_length:
            state = env.reset()
            writer.add_scalar("evaluator/episode_reward", episode_reward,
                              episode_idx)
            writer.add_scalar("evaluator/episode_length", episode_length,
                              episode_idx)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1
            param = param_queue.get()
            model.load_state_dict(param)
            print("Updated Parameter..")
Пример #9
0
def vector_exploration(args, actor_id, n_actors, replay_ip, param_queue):
    ctx = zmq.Context()
    batch_socket = ctx.socket(zmq.DEALER)
    batch_socket.setsockopt(zmq.IDENTITY,
                            pickle.dumps('actor-{}'.format(actor_id)))
    batch_socket.connect('tcp://{}:51001'.format(replay_ip))
    outstanding = 0

    writer = SummaryWriter(comment="-{}-actor{}".format(args.env, actor_id))

    num_envs = args.num_envs_per_worker
    envs = [
        wrap_atari_dqn(make_atari(args.env), args) for _ in range(num_envs)
    ]

    if args.seed is not None:
        seeds = args.seed + actor_id * num_envs + np.arange(num_envs)
        utils.set_global_seeds(seeds[0], use_torch=True)
        for seed, env in zip(seeds, envs):
            env.seed(int(seed))

    model = DuelingDQN(envs[0], args)
    model = torch.jit.trace(model, torch.zeros((1, 4, 84, 84)))
    _actor_id = np.arange(num_envs) + actor_id * num_envs
    n_actors = n_actors * num_envs
    epsilons = args.eps_base**(1 + _actor_id / (n_actors - 1) * args.eps_alpha)
    storages = [
        BatchStorage(args.n_steps, args.gamma) for _ in range(num_envs)
    ]

    param = param_queue.get(block=True)
    model.load_state_dict(param)
    param = None
    print("%d: Received First Parameter!" % actor_id)

    actor_idx = 0
    tb_idx = 0
    episode_rewards = np.array([0] * num_envs)
    episode_lengths = np.array([0] * num_envs)
    states = np.array([env.reset() for env in envs])
    tb_dict = {key: [] for key in ['episode_reward', 'episode_length']}
    step_t = time.time()
    ref_t = 0
    sim_t = 0
    while True:
        if actor_idx * num_envs * n_actors <= args.initial_exploration_samples:  # initial random exploration
            random_idx = np.arange(num_envs)
        else:
            random_idx, = np.where(np.random.random(num_envs) <= epsilons)
        _t = time.time()
        with torch.no_grad():
            states_tensor = torch.tensor(states, dtype=torch.float32)
            q_values = model(states_tensor).detach().numpy()
        ref_t += time.time() - _t
        actions = np.argmax(q_values, 1)
        actions[random_idx] = np.random.choice(envs[0].action_space.n,
                                               len(random_idx))

        for i, (state, q_value, action, env, storage) in enumerate(
                zip(states, q_values, actions, envs, storages)):
            _t = time.time()
            next_state, reward, done, _ = env.step(action)
            sim_t += time.time() - _t
            storage.add(np.array(state), reward, action, done, q_value, _t,
                        episode_lengths[i])
            states[i] = next_state
            episode_rewards[i] += reward
            episode_lengths[i] += 1
            if done or episode_lengths[i] == args.max_episode_length:
                states[i] = env.reset()
                tb_idx += 1
                tb_dict["episode_reward"].append(episode_rewards[i])
                tb_dict["episode_length"].append(episode_lengths[i])
                episode_rewards[i] = 0
                episode_lengths[i] = 0
                if tb_idx % args.tb_interval == 0:
                    writer.add_scalar('actor/episode_reward_mean',
                                      np.mean(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/episode_reward_max',
                                      np.max(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/episode_reward_min',
                                      np.min(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/episode_reward_std',
                                      np.std(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/episode_length_mean',
                                      np.mean(tb_dict['episode_length']),
                                      tb_idx)
                    tb_dict['episode_reward'].clear()
                    tb_dict['episode_length'].clear()
                    writer.add_scalar('actor/step_time',
                                      (time.time() - step_t) /
                                      np.sum(episode_lengths), tb_idx)
                    writer.add_scalar('actor/step_inference_time',
                                      ref_t / np.sum(episode_lengths), tb_idx)
                    writer.add_scalar('actor/step_simulation_time',
                                      sim_t / np.sum(episode_lengths), tb_idx)
                    ref_t = 0
                    sim_t = 0
                    step_t = time.time()

        actor_idx += 1

        if actor_idx % args.update_interval == 0:
            try:
                param = param_queue.get(block=False)
                model.load_state_dict(param)
                print("%d: Updated Parameter.." % actor_id)
            except queue.Empty:
                pass

        if sum(len(storage)
               for storage in storages) >= args.send_interval * num_envs:
            batch = []
            prios = []
            for storage in storages:
                _batch, _prios = storage.make_batch()
                batch.append(_batch)
                prios.append(_prios)
                storage.reset()
            batch = [np.concatenate(v) for v in zip(*batch)]
            prios = np.concatenate(prios)
            data = pickle.dumps((batch, prios))
            batch, prios = None, None
            while outstanding >= args.max_outstanding:
                batch_socket.recv()
                outstanding -= 1
            batch_socket.send(data, copy=False)
            outstanding += 1
class Agent:
    """Interacts with and learns from the environment."""
    def __init__(self, config):
        """Initialize an Agent object"""
        self.seed = random.seed(config["general"]["seed"])
        self.config = config

        # Q-Network
        self.q = DuelingDQN(config).to(DEVICE)
        self.q_target = DuelingDQN(config).to(DEVICE)

        self.optimizer = optim.RMSprop(self.q.parameters(),
                                       lr=config["agent"]["learning_rate"])
        self.criterion = F.mse_loss

        self.memory = ReplayBuffer(config)
        # Initialize time step (for updating every UPDATE_EVERY steps)
        self.t_step = 0

    def save_experiences(self, state, action, reward, next_state, done):
        """Prepare and save experience in replay memory"""
        reward = np.clip(reward, -1.0, 1.0)
        self.memory.add(state, action, reward, next_state, done)

    def _current_step_is_a_learning_step(self):
        """Check if the current step is an update step"""
        self.t_step = (self.t_step + 1) % self.config["agent"]["update_rate"]
        return self.t_step == 0

    def _enough_samples_in_memory(self):
        """Check if minimum amount of samples are in memory"""
        return len(self.memory) > self.config["train"]["batch_size"]

    def epsilon_greedy_action_selection(self, action_values, eps):
        """Epsilon-greedy action selection"""
        if random.random() > eps:
            return np.argmax(action_values.cpu().data.numpy())
        else:
            return random.choice(
                np.arange(self.config["general"]["action_size"]))

    def act(self, state, eps=0.0):
        """Returns actions for given state as per current policy"""
        state = torch.from_numpy(state).float().unsqueeze(0).to(DEVICE)
        self.q.eval()
        with torch.no_grad():
            action_values = self.q(state)
        self.q.train()

        return self.epsilon_greedy_action_selection(action_values, eps)

    def _calc_loss(self, states, actions, rewards, next_states, dones):
        """Calculates loss for a given experience batch"""
        q_eval = self.q(states).gather(1, actions)
        q_eval_next = self.q(next_states)
        _, q_argmax = q_eval_next.detach().max(1)
        q_next = self.q_target(next_states)
        q_next = q_next.gather(1, q_argmax.unsqueeze(1))
        q_target = rewards + (self.config["agent"]["gamma"] * q_next *
                              (1 - dones))
        loss = self.criterion(q_eval, q_target)
        return loss

    def _update_weights(self, loss):
        """update the q network weights"""
        torch.nn.utils.clip_grad.clip_grad_value_(self.q.parameters(), 1.0)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def learn(self):
        """Update network using one sample of experience from memory"""
        if self._current_step_is_a_learning_step(
        ) and self._enough_samples_in_memory():
            states, actions, rewards, next_states, dones = self.memory.sample(
                self.config["train"]["batch_size"])
            loss = self._calc_loss(states, actions, rewards, next_states,
                                   dones)
            self._update_weights(loss)
            self._soft_update(self.q, self.q_target)

    def _soft_update(self, local_model, target_model):
        """Soft update target network parameters: θ_target = τ*θ_local + (1 - τ)*θ_target"""
        for target_param, local_param in zip(target_model.parameters(),
                                             local_model.parameters()):
            target_param.data.copy_(
                self.config["agent"]["tau"] * local_param.data +
                (1.0 - self.config["agent"]["tau"]) * target_param.data)

    def save(self):
        """Save the network weights"""
        helper.mkdir(
            os.path.join(".", *self.config["general"]["checkpoint_dir"],
                         self.config["general"]["env_name"]))
        current_date_time = helper.get_current_date_time()
        current_date_time = current_date_time.replace(" ", "__").replace(
            "/", "_").replace(":", "_")

        torch.save(
            self.q.state_dict(),
            os.path.join(".", *self.config["general"]["checkpoint_dir"],
                         self.config["general"]["env_name"],
                         "ckpt_" + current_date_time))

    def load(self):
        """Load latest available network weights"""
        list_of_files = glob.glob(
            os.path.join(".", *self.config["general"]["checkpoint_dir"],
                         self.config["general"]["env_name"], "*"))
        latest_file = max(list_of_files, key=os.path.getctime)
        self.q.load_state_dict(torch.load(latest_file))
        self.q_target.load_state_dict(torch.load(latest_file))
Пример #11
0
class train_DQN():
    def __init__(self,
                 env_id,
                 seed=0,
                 lr=1e-5,
                 n_step=3,
                 gamma=0.99,
                 n_workers=20,
                 max_norm=40,
                 target_update_interval=2500,
                 save_interval=5000,
                 batch_size=64,
                 buffer_size=1e6,
                 prior_alpha=0.6,
                 prior_beta=0.4,
                 publish_param_interval=32,
                 max_step=1e5):
        self.env = gym.make(env_id)
        self.seed = seed
        self.lr = lr
        self.n_step = n_step
        self.gamma = gamma
        self.max_norm = max_norm
        self.target_update_interval = target_update_interval
        self.save_interval = save_interval
        self.publish_param_interval = publish_param_interval
        self.batch_size = batch_size
        self.prior_beta = prior_beta
        self.max_step = max_step

        self.buffer = CustomPrioritizedReplayBuffer(size=buffer_size,
                                                    alpha=prior_alpha)
        self.device = torch.device(
            "cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = DuelingDQN(self.env).to(self.device)
        self.tgt_model = DuelingDQN(self.env).to(self.device)
        self.tgt_model.load_state_dict(self.model.state_dict())
        self.optimizer = torch.optim.RMSprop(self.model.parameters(),
                                             self.lr,
                                             alpha=0.95,
                                             eps=1.5e-7,
                                             centered=True)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer,
                                                         step_size=1000,
                                                         gamma=0.99)
        self.beta_by_frame = lambda frame_idx: min(
            1.0, self.prior_beta + frame_idx * (1.0 - self.prior_beta) / 1000)
        self.batch_recorder = BatchRecorder(env_id=env_id,
                                            env_seed=seed,
                                            n_workers=n_workers,
                                            buffer=self.buffer,
                                            n_steps=n_step,
                                            gamma=gamma,
                                            max_episode_length=50000)
        self.writer = SummaryWriter(
            comment="-{}-learner".format(self.env.unwrapped.spec.id))

    def train(self):
        utils.set_global_seeds(self.seed, use_torch=True)

        learn_idx = 0
        while True:
            beta = self.beta_by_frame(learn_idx)
            states, actions, rewards, next_states, dones, weights, idxes = self.buffer.sample(
                self.batch_size, beta)
            states = torch.FloatTensor(states).to(self.device)
            actions = torch.LongTensor(actions).to(self.device)
            rewards = torch.FloatTensor(rewards).to(self.device)
            next_states = torch.FloatTensor(next_states).to(self.device)
            dones = torch.FloatTensor(dones).to(self.device)
            weights = torch.FloatTensor(weights).to(self.device)
            batch = (states, actions, rewards, next_states, dones, weights)

            loss, prios = utils.compute_loss(self.model, self.tgt_model, batch,
                                             self.n_step, self.gamma)

            self.scheduler.step()
            grad_norm = utils.update_parameters(loss, self.model,
                                                self.optimizer, self.max_norm)

            self.buffer.update_priorities(idxes, prios)

            batch, idxes, prios = None, None, None
            learn_idx += 1

            self.writer.add_scalar("learner/loss", loss, learn_idx)
            self.writer.add_scalar("learner/grad_norm", grad_norm, learn_idx)

            if learn_idx % self.target_update_interval == 0:
                print("Updating Target Network..")
                self.tgt_model.load_state_dict(self.model.state_dict())
            if learn_idx % self.save_interval == 0:
                print("Saving Model..")
                torch.save(self.model.state_dict(),
                           "model{}.pth".format(learn_idx))
            if learn_idx % self.publish_param_interval == 0:
                self.batch_recorder.set_worker_weights(
                    copy.deepcopy(self.model))
            if learn_idx >= self.max_step:
                torch.save(self.model.state_dict(),
                           "model{}.pth".format(learn_idx))
                self.batch_recorder.cleanup()
                break

    def load_model(self, idx):
        with open("model{}.pth".format(idx), "rb") as f:
            print("loading weights_{}".format(idx))
            self.model.load_state_dict(torch.load(f, map_location="cpu"))

    def sampling_data(self):
        self.batch_recorder.record_batch()
Пример #12
0
def exploration(args, actor_id, n_actors, replay_ip, param_queue):
    ctx = zmq.Context()
    batch_socket = ctx.socket(zmq.DEALER)
    batch_socket.setsockopt(zmq.IDENTITY, pickle.dumps('actor-{}'.format(actor_id)))
    batch_socket.connect('tcp://{}:51001'.format(replay_ip))
    outstanding = 0

    writer = SummaryWriter(comment="-{}-actor{}".format(args.env, actor_id))
    
    env = RunTagEnv(width=5, height=5, number_of_subordinates=1, max_steps=1000)

    # env = make_atari(args.env)
    # env = wrap_atari_dqn(env, args)

    seed = args.seed + actor_id
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    blue_commander_model = DuelingDQN(env)
    red_commander_model = RandomPolicy(env)

    blue_commander_policy = CommanderPolicy(blue_commander_model)
    red_commander_policy = CommanderPolicy(red_commander_model)
    subordinate_policy = SubordinatePolicy()

    # model = DuelingDQN(env)
    epsilon = args.eps_base ** (1 + actor_id / (n_actors - 1) * args.eps_alpha)
    storage = BatchStorage(args.n_steps, args.gamma)

    param = param_queue.get(block=True)
    blue_commander_model.load_state_dict(param)
    param = None
    print("Received First Parameter!")

    episode_reward, episode_length, episode_idx, actor_idx = 0, 0, 0, 0
    states = env.reset()
    while True:
        #action, q_values = model.act(torch.FloatTensor(np.array(state[0])), epsilon)
        state = flatten(env, states[0])
        blue_com_action, q_values = blue_commander_policy.next_action(state, epsilon)
        blue_sub_action = subordinate_policy.next_action(states[1])
        red_com_action = red_commander_policy.next_action(flatten(env, states[2]))
        red_sub_action = subordinate_policy.next_action(states[3])
        actions = [
                blue_com_action,
                blue_sub_action,
                red_com_action,
                red_sub_action,
            ]


        next_states, rewards, dones, _ = env.step(actions)
        next_state, reward, done, action = next_states[0], rewards[0], dones[0], actions[0]
        # print('state:', state)
        storage.add(state, reward, action, done, q_values)

        states = next_states
        episode_reward += reward
        episode_length += 1
        actor_idx += 1

        if done or episode_length == args.max_episode_length:
            state = env.reset()
            writer.add_scalar("actor/episode_reward", episode_reward, episode_idx)
            writer.add_scalar("actor/episode_length", episode_length, episode_idx)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1

        if actor_idx % args.update_interval == 0:
            try:
                param = param_queue.get(block=False)
                model.load_state_dict(param)
                print("Updated Parameter..")
            except queue.Empty:
                pass

        if len(storage) == args.send_interval:
            batch, prios = storage.make_batch()
            data = pickle.dumps((batch, prios))
            batch, prios = None, None
            storage.reset()
            while outstanding >= args.max_outstanding:
                batch_socket.recv()
                outstanding -= 1
            batch_socket.send(data, copy=False)
            outstanding += 1
            print("Sending Batch..")
Пример #13
0
def evaluator(args):
    comm = global_dict['comm_world']
    writer = SummaryWriter(log_dir=os.path.join(args['log_dir'], 'eval'))

    args['clip_rewards'] = False
    args['episode_life'] = False
    env = make_atari(args['env'])
    env = wrap_atari_dqn(env, args)

    seed = args['seed'] - 1
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    torch.set_num_threads(1)
    model = DuelingDQN(env, args)

    recv_param_buf = bytearray(100 * 1024 * 1024)
    comm.Send(b'', dest=global_dict['rank_learner'])
    comm.Recv(buf=recv_param_buf, source=global_dict['rank_learner'])
    param = pickle.loads(recv_param_buf)
    model.load_state_dict(param)

    episode_reward, episode_length, episode_idx = 0, 0, 0
    state = env.reset()
    tb_dict = {k: [] for k in ['episode_reward', 'episode_length']}
    while True:
        action, _ = model.act(torch.FloatTensor(np.array(state)), 0.)
        next_state, reward, done, _ = env.step(action)

        state = next_state
        episode_reward += reward
        episode_length += 1

        if done or episode_length == args['max_episode_length']:
            state = env.reset()
            tb_dict["episode_reward"].append(episode_reward)
            tb_dict["episode_length"].append(episode_length)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1
            comm.Send(b'', dest=global_dict['rank_learner'])
            comm.Recv(buf=recv_param_buf, source=global_dict['rank_learner'])
            param = pickle.loads(recv_param_buf)
            model.load_state_dict(param)

            if (episode_idx *
                    args['num_envs_per_worker']) % args['tb_interval'] == 0:
                writer.add_scalar('evaluator/1_episode_reward_mean',
                                  np.mean(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/2_episode_reward_max',
                                  np.max(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/3_episode_reward_min',
                                  np.min(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/4_episode_reward_std',
                                  np.std(tb_dict['episode_reward']),
                                  episode_idx)
                writer.add_scalar('evaluator/5_episode_length_mean',
                                  np.mean(tb_dict['episode_length']),
                                  episode_idx)
                tb_dict['episode_reward'].clear()
                tb_dict['episode_length'].clear()
Пример #14
0
class Worker(mp.Process):
    def __init__(self, worker_id, env_id, seed, epsilon, size, lock,
                n_steps, gamma, send_interval, task_queue, buffer, max_episode_length):
        mp.Process.__init__(self)
        self.worker_id = worker_id
        self.env = gym.make(env_id)
        self.seed = seed
        self.task_queue = task_queue
        self.buffer = buffer
        self.max_episode_length = max_episode_length
        self.send_interval = send_interval
        self.storage = BatchStorage(n_steps, gamma)
        self.size = size
        self.memory = []
        
        self.model = DuelingDQN(self.env)
        # self.writer = SummaryWriter(comment="-{}-actor{}".format(env_id, worker_id))
        self.lock = lock

        self.set_all_seeds()
        self.epsilon = epsilon
    def set_all_seeds(self):
        self.env.seed(self.seed)
        np.random.seed(self.seed)
        random.seed(self.seed)
        torch.manual_seed(self.seed)
        torch.cuda.manual_seed_all(self.seed)
    def update_weights(self, DQN_state_dict):
        self.model.load_state_dict(DQN_state_dict)
    def record_batch(self):

        episode_reward, episode_length, episode_idx, actor_idx = 0, 0, 0, 0
        state = self.env.reset()
        self.storage.reset()
        self.memory = []
        # while actor_idx < self.size:
        while True:
            action, q_values = self.model.act(torch.FloatTensor(np.array(state)), self.epsilon)
            next_state, reward, done, _ = self.env.step(action)
            
            
            self.storage.add(state, reward, action, done, q_values)

            state = next_state
            episode_reward += reward
            episode_length += 1
            actor_idx += 1
            if done or episode_length >= self.max_episode_length:
                state = self.env.reset()
                # self.writer.add_scalar("actor/episode_reward", episode_reward, episode_idx)
                # self.writer.add_scalar("actor/episode_length", episode_length, episode_idx)
                episode_reward = 0
                episode_length = 0
                episode_idx += 1


            if done or len(self.storage) == self.send_interval:
                batch, prios = self.storage.make_batch()
                self.memory.append((*batch, prios))
                # for i in range(len(prios)):
                #     self.buffer.add(batch[0][i],batch[1][i],batch[2][i],batch[3][i],batch[4][i],prios[i])

                batch, prios = None, None
                self.storage.reset()
            if done:
                break
    def run(self):
        while True:
            ########## run loop
            task = self.task_queue.get(block=True)
            if task["desc"] == "record_batch":
                # print("start record batch")
                self.record_batch()
                self.buffer.put(self.memory)
                self.task_queue.task_done()
                # print("record batch done")
            elif task["desc"] == "set_pi_weights":
                # print("set weight")
                self.update_weights(task["pi_state_dict"])
                self.task_queue.task_done()
                # print("set weight done")
            elif task["desc"] == "cleanup":
                # print("clean up")
                self.env.close()
                self.task_queue.task_done()
Пример #15
0
def exploration(args, actor_id, n_actors, replay_ip, param_queue, sample_enque, sample_deque):
    ctx = zmq.Context()
    batch_socket = ctx.socket(zmq.DEALER)
    batch_socket.setsockopt(zmq.IDENTITY, pickle.dumps('actor-{}'.format(actor_id)))
    batch_socket.connect('tcp://{}:51001'.format(replay_ip))

    writer = SummaryWriter(comment="-{}-actor{}".format(args.env, actor_id))

    env = make_atari(args.env)
    env = wrap_atari_dqn(env, args)

    seed = args.seed + actor_id
    utils.set_global_seeds(seed, use_torch=True)
    env.seed(seed)

    model = DuelingDQN(env)
    epsilon = args.eps_base ** (1 + actor_id / (n_actors - 1) * args.eps_alpha)
    storage = BatchStorage(args.n_steps, args.gamma)

    param = param_queue.get(block=True)
    model.load_state_dict(param)
    param = None
    print("Received First Parameter!")

    episode_reward, episode_length, episode_idx, actor_idx = 0, 0, 0, 0
    state = env.reset()
    while True:
        action, q_values = model.act(torch.FloatTensor(np.array(state)), epsilon)
        next_state, reward, done, _ = env.step(action)
        storage.add(state, reward, action, done, q_values)

        state = next_state
        episode_reward += reward
        episode_length += 1
        actor_idx += 1

        if done or episode_length == args.max_episode_length:
            state = env.reset()
            writer.add_scalar("actor/episode_reward", episode_reward, episode_idx)
            writer.add_scalar("actor/episode_length", episode_length, episode_idx)
            episode_reward = 0
            episode_length = 0
            episode_idx += 1

        if actor_idx % args.update_interval == 0:
            try:
                param = param_queue.get(block=False)
                model.load_state_dict(param)
                print("Updated Parameter..")
            except queue.Empty:
                pass

        # get sample batch after each step
        while sample_enque:
            idxes = sample_enque.get()
            sample_deque.put(storage.get_sample_batch(idxes))

        # only pass the prios and get indexes from ReplayBuffer
        if len(storage) == args.send_interval:
            batch, prios = storage.make_batch()
            data = pickle.dumps(prios)
            batch, prios = None, None
            storage.reset()
            batch_socket.send(data, copy=False)
            _, idxes = batch_socket.recv_multipart(copy=False)
            storage.add_batch(batch, idxes)
Пример #16
0
class train_DQN():
    def __init__(self, env_id, max_step = 1e5, prior_alpha = 0.6, prior_beta_start = 0.4, 
                    epsilon_start = 1.0, epsilon_final = 0.01, epsilon_decay = 500,
                    batch_size = 32, gamma = 0.99, target_update_interval=1000, save_interval = 1e4,
                    ):
        self.prior_beta_start = prior_beta_start
        self.max_step = int(max_step)
        self.batch_size = batch_size
        self.gamma = gamma
        self.target_update_interval = target_update_interval
        self.save_interval = save_interval


        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.env = gym.make(env_id)
        self.model = DuelingDQN(self.env).to(self.device)
        self.target_model = DuelingDQN(self.env).to(self.device)
        self.target_model.load_state_dict(self.model.state_dict())
        self.replay_buffer = PrioritizedReplayBuffer(100000,alpha=prior_alpha)
        self.optimizer = optim.Adam(self.model.parameters())
        self.writer = SummaryWriter(comment="-{}-learner".format(self.env.unwrapped.spec.id))


        # decay function
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer,step_size=1000,gamma=0.99)
        self.beta_by_frame = lambda frame_idx: min(1.0, self.prior_beta_start + frame_idx * (1.0 - self.prior_beta_start) / 1000)
        self.epsilon_by_frame = lambda frame_idx: epsilon_final + (epsilon_start - epsilon_final) * math.exp(-1. * frame_idx / epsilon_decay)
        
    def update_target(self,current_model, target_model):
        target_model.load_state_dict(current_model.state_dict())
    def compute_td_loss(self,batch_size, beta):
        state, action, reward, next_state, done, weights, indices  = self.replay_buffer.sample(batch_size, beta) 

        state      = torch.FloatTensor(state).to(self.device)
        next_state = torch.FloatTensor(next_state).to(self.device)
        action     = torch.LongTensor(action).to(self.device)
        reward     = torch.FloatTensor(reward).to(self.device)
        done       = torch.FloatTensor(done).to(self.device)
        weights    = torch.FloatTensor(weights).to(self.device)
        batch = (state, action, reward, next_state, done, weights)

        # q_values      = self.model(state)
        # next_q_values = self.target_model(next_state)

        # q_value          = q_values.gather(1, action.unsqueeze(1)).squeeze(1)
        # next_q_value     = next_q_values.max(1)[0]
        # expected_q_value = reward + self.gamma * next_q_value * (1 - done)
        
        # td_error = torch.abs(expected_q_value.detach() - q_value)
        # loss  = (td_error).pow(2) * weights
        # prios = loss+1e-5#0.9 * torch.max(td_error)+(1-0.9)*td_error
        # loss  = loss.mean()
        loss, prios = utils.compute_loss(self.model,self.target_model, batch,1)
            
        self.optimizer.zero_grad()
        loss.backward()
        self.scheduler.step()
        self.replay_buffer.update_priorities(indices, prios)
        self.optimizer.step()
        return loss    
    def train(self):
        losses = []
        all_rewards = []
        episode_reward = 0
        episode_idx = 0
        episode_length = 0
        state = self.env.reset()
        for frame_idx in range(self.max_step):
            epsilon = self.epsilon_by_frame(frame_idx)
            action,_ = self.model.act(torch.FloatTensor((state)).to(self.device), epsilon)
            next_state, reward, done, _ = self.env.step(action)
            self.replay_buffer.add(state, action, reward, next_state, done)
            
            state = next_state
            episode_reward += reward
            
            episode_length += 1
            if done:
                state = self.env.reset()
                all_rewards.append(episode_reward)
                self.writer.add_scalar("actor/episode_reward", episode_reward, episode_idx)
                self.writer.add_scalar("actor/episode_length", episode_length, episode_idx)
                # print("episode: ",episode_idx, " reward: ", episode_reward)
                episode_reward = 0
                episode_length = 0
                episode_idx += 1
                
            if len(self.replay_buffer) > self.batch_size:
                beta = self.beta_by_frame(frame_idx)
                loss = self.compute_td_loss(self.batch_size, beta)
                losses.append(loss.item())
                self.writer.add_scalar("learner/loss", loss, frame_idx)
                
            if frame_idx % self.target_update_interval == 0:
                print("update target...")
                self.update_target(self.model, self.target_model)

            if frame_idx % self.save_interval == 0 or frame_idx == self.max_step-1:
                print("save model...")
                self.save_model(frame_idx)

            
    def save_model(self, idx):
        torch.save(self.model.state_dict(), "./model{}.pth".format(idx))
    def load_model(self,idx):
         with open("model{}.pth".format(idx), "rb") as f:
                print("loading weights_{}".format(idx))
                self.model.load_state_dict(torch.load(f,map_location="cpu"))
Пример #17
0
def actor(args, actor_id):
    comm = global_dict['comm_local']
    writer = SummaryWriter(log_dir=os.path.join(
        args['log_dir'], f'{global_dict["unit_idx"]}-actor{actor_id}'))

    num_envs = args['num_envs_per_worker']
    envs = [
        wrap_atari_dqn(make_atari(args['env']), args) for _ in range(num_envs)
    ]

    if args['seed'] is not None:
        seeds = args['seed'] + actor_id * num_envs + np.arange(num_envs)
        utils.set_global_seeds(seeds[0], use_torch=True)
        for seed, env in zip(seeds, envs):
            env.seed(int(seed))

    model = DuelingDQN(envs[0], args)
    model = torch.jit.trace(model, torch.zeros((1, 4, 84, 84)))
    _actor_id = (np.arange(num_envs) + actor_id *
                 num_envs) * args['num_units'] + global_dict['unit_idx']
    n_actors = args['num_actors'] * num_envs * args['num_units']
    epsilons = args['eps_base']**(1 + _actor_id /
                                  (n_actors - 1) * args['eps_alpha'])
    storages = [
        BatchStorage(args['n_steps'], args['gamma']) for _ in range(num_envs)
    ]

    recv_param_buf = bytearray(100 * 1024 * 1024)
    recv_param_request = None
    send_batch_request = None

    actor_idx = 0
    tb_idx = 0
    episode_rewards = np.array([0] * num_envs)
    episode_lengths = np.array([0] * num_envs)
    states = np.array([env.reset() for env in envs])
    tb_dict = {
        key: []
        for key in
        ['episode_reward', 'episode_length', 'kept_sample_percentage']
    }
    step_t = time.time()
    inf_t = 0
    sim_t = 0

    def make_episilons():
        return epsilons

    while True:
        if recv_param_request and recv_param_request.Test():
            param = pickle.loads(recv_param_buf)
            model.load_state_dict(param)
            recv_param_request = None
        if actor_idx * num_envs * n_actors <= args[
                'initial_exploration_samples']:  # initial random exploration
            random_idx = np.arange(num_envs)
        else:
            random_idx, = np.where(
                np.random.random(num_envs) <= make_episilons())
        _t = time.time()
        with torch.no_grad():
            states_tensor = torch.tensor(states, dtype=torch.float32)
            q_values = model(states_tensor).detach().numpy()
        inf_t += time.time() - _t
        actions = np.argmax(q_values, 1)
        actions[random_idx] = np.random.choice(envs[0].action_space.n,
                                               len(random_idx))

        for i, (state, q_value, action, env, storage) in enumerate(
                zip(states, q_values, actions, envs, storages)):
            _t = time.time()
            next_state, reward, done, _ = env.step(action)
            try:
                real_done = env.was_real_done
            except:
                real_done = done
            sim_t += time.time() - _t
            storage.add(np.array(state), reward, action, done, real_done,
                        q_value, _t, episode_lengths[i])
            states[i] = next_state
            episode_rewards[i] += reward
            episode_lengths[i] += 1
            if done or episode_lengths[i] == args['max_episode_length']:
                states[i] = env.reset()
            if real_done or episode_lengths[i] == args['max_episode_length']:
                tb_idx += 1
                tb_dict["episode_reward"].append(episode_rewards[i])
                tb_dict["episode_length"].append(episode_lengths[i])
                episode_rewards[i] = 0
                episode_lengths[i] = 0
                if tb_idx % args['tb_interval'] == 0:
                    writer.add_scalar('actor/1_episode_reward_mean',
                                      np.mean(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/2_episode_reward_max',
                                      np.max(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/3_episode_reward_min',
                                      np.min(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/4_episode_reward_std',
                                      np.std(tb_dict['episode_reward']),
                                      tb_idx)
                    writer.add_scalar('actor/5_episode_length_mean',
                                      np.mean(tb_dict['episode_length']),
                                      tb_idx)
                    tb_dict['episode_reward'].clear()
                    tb_dict['episode_length'].clear()
                    writer.add_scalar('actor/6_step_time',
                                      (time.time() - step_t) /
                                      np.sum(episode_lengths), tb_idx)
                    writer.add_scalar('actor/7_step_inference_time',
                                      inf_t / np.sum(episode_lengths), tb_idx)
                    writer.add_scalar('actor/8_step_simulation_time',
                                      sim_t / np.sum(episode_lengths), tb_idx)
                    writer.add_scalar(
                        'actor/9_kept_sample_percentage',
                        np.mean(tb_dict['kept_sample_percentage']), tb_idx)
                    inf_t = 0
                    sim_t = 0
                    step_t = time.time()
                    tb_dict['kept_sample_percentage'].clear()

        actor_idx += 1

        if actor_idx % args['update_interval'] == 0:
            if recv_param_request is not None:
                print(
                    f"actor {global_dict['unit_idx']}-{actor_id}: last recv param request is not complete!"
                )
                sys.stdout.flush()
            else:
                comm.Send(b'', dest=global_dict['rank_learner'])
                recv_param_request = comm.Irecv(
                    buf=recv_param_buf, source=global_dict['rank_learner'])

        if sum(len(storage)
               for storage in storages) >= args['send_interval'] * num_envs:
            batch = []
            prios = []
            for storage in storages:
                _batch, _prios = storage.make_batch()
                batch.append(_batch)
                prios.append(_prios)
                storage.reset()
            batch = [np.concatenate(v) for v in zip(*batch)]
            prios = np.concatenate(prios)
            threshold = args['sample_filter_threshold']
            prios_mask = prios > np.max(prios) * threshold
            tb_dict['kept_sample_percentage'].append(
                np.sum(prios_mask) / len(prios_mask))
            prios = prios[prios_mask]
            batch = [i[prios_mask] for i in batch]
            data = pickle.dumps((batch, prios))
            if send_batch_request is not None:
                send_batch_request.wait()
            send_batch_request = comm.Isend(data,
                                            dest=global_dict['rank_replay'],
                                            tag=utils.TAG_RECV_BATCH)
Пример #18
0
def learner(args):
    comm_cross = global_dict['comm_cross']
    hvd.init(comm=comm_cross)
    torch.cuda.set_device(hvd.local_rank())
    env = wrap_atari_dqn(make_atari(args['env']), args)
    # utils.set_global_seeds(args['seed'], use_torch=True)

    device = args['device']
    model = DuelingDQN(env, args).to(device)
    if os.path.exists('model.pth'):
        # model.load_state_dict(torch.load('model.pth'))
        pass

    tgt_model = DuelingDQN(env, args).to(device)
    del env

    writer = SummaryWriter(log_dir=os.path.join(
        args['log_dir'], f'{global_dict["unit_idx"]}-learner'))
    # optimizer = torch.optim.SGD(model.parameters(), 1e-5 * args['num_units'], momentum=0.8)
    # optimizer = torch.optim.RMSprop(model.parameters(), args['lr'], alpha=0.95, eps=1.5e-7, centered=True)
    optimizer = torch.optim.Adam(model.parameters(),
                                 args['lr'] * args['num_units'])
    optimizer = hvd.DistributedOptimizer(
        optimizer, named_parameters=model.named_parameters())
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    tgt_model.load_state_dict(model.state_dict())
    if args['dynamic_gradient_clip']:
        grad_norm_running_mean = args['gradient_norm_running_mean']
        grad_norm_lambda = args['gradient_norm_lambda']

    batch_queue = queue.Queue(maxsize=3)
    prios_queue = queue.Queue(maxsize=4)
    param_queue = queue.Queue(maxsize=3)
    threading.Thread(target=recv_batch, args=(batch_queue, )).start()
    threading.Thread(target=send_prios, args=(prios_queue, )).start()
    threading.Thread(target=send_param, args=(param_queue, )).start()
    if global_dict['unit_idx'] == 0:
        threading.Thread(target=send_param_evaluator,
                         args=(param_queue, )).start()

    prefetcher = data_prefetcher(batch_queue, args['cuda'])

    learn_idx = 0
    ts = time.time()
    tb_dict = {
        k: []
        for k in [
            'loss', 'grad_norm', 'max_q', 'mean_q', 'min_q',
            'batch_queue_size', 'prios_queue_size'
        ]
    }
    first_rount = True
    while True:
        (*batch, idxes) = prefetcher.next()
        if first_rount:
            print("start training")
            sys.stdout.flush()
            first_rount = False
        loss, prios, q_values = utils.compute_loss(model, tgt_model, batch,
                                                   args['n_steps'],
                                                   args['gamma'])

        optimizer.zero_grad()
        loss.backward()
        if args['dynamic_gradient_clip']:
            grad_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(),
                grad_norm_running_mean * args['clipping_threshold'])
            grad_norm_running_mean = grad_norm_running_mean * grad_norm_lambda + \
                min(grad_norm, grad_norm_running_mean * args['clipping_threshold']) * (1-grad_norm_lambda)
        else:
            grad_norm = torch.norm(
                torch.stack([
                    torch.norm(p.grad.detach(), 2) for p in model.parameters()
                ]), 2)
        # global_prios_sum = np.array(prios_sum)
        # comm_cross.Allreduce(MPI.IN_PLACE, global_prios_sum.data)
        # global_prios_sum = float(global_prios_sum)
        # scale = prios_sum / global_prios_sum
        if args['dynamic_gradient_clip'] and args[
                'dropping_threshold'] and grad_norm > grad_norm_running_mean * args[
                    'dropping_threshold']:
            pass
        else:
            optimizer.step()

        prios_queue.put((idxes, prios))
        learn_idx += 1
        tb_dict["loss"].append(float(loss))
        tb_dict["grad_norm"].append(float(grad_norm))
        tb_dict["max_q"].append(float(torch.max(q_values)))
        tb_dict["mean_q"].append(float(torch.mean(q_values)))
        tb_dict["min_q"].append(float(torch.min(q_values)))
        tb_dict["batch_queue_size"].append(batch_queue.qsize())
        tb_dict["prios_queue_size"].append(prios_queue.qsize())

        if learn_idx % args['target_update_interval'] == 0:
            tgt_model.load_state_dict(model.state_dict())
        if learn_idx % args['save_interval'] == 0 and global_dict[
                'unit_idx'] == 0:
            torch.save(model.state_dict(), "model.pth")
        if learn_idx % args['publish_param_interval'] == 0:
            param_queue.put(model.state_dict())
        if learn_idx % args['tb_interval'] == 0:
            bps = args['tb_interval'] / (time.time() - ts)
            for i, (k, v) in enumerate(tb_dict.items()):
                writer.add_scalar(f'learner/{i+1}_{k}', np.mean(v), learn_idx)
                v.clear()
            writer.add_scalar(f"learner/{i+2}_BPS", bps, learn_idx)
            ts = time.time()
Пример #19
0
def train(args, n_actors, batch_queue, prios_queue, param_queue):
    """
    thread to fill parameter queue
    """
    def _fill_param():
        while True:
            model_dict = {}
            state_dict = model.state_dict()
            for k, v in state_dict.items():
                model_dict[k] = v.cpu().numpy()
            param_queue.put(model_dict)

    env = wrapper.make_atari(args.env)
    env = wrapper.wrap_atari_dqn(env, args)
    utils.set_global_seeds(args.seed, use_torch=True)

    model = DuelingDQN(env).to(args.device)
    tgt_model = DuelingDQN(env).to(args.device)
    tgt_model.load_state_dict(model.state_dict())

    writer = SummaryWriter(comment="-{}-learner".format(args.env))
    # optimizer = torch.optim.Adam(model.parameters(), args.lr)
    optimizer = torch.optim.RMSprop(model.parameters(),
                                    args.lr,
                                    alpha=0.95,
                                    eps=1.5e-7,
                                    centered=True)
    model_dict = {}
    state_dict = model.state_dict()
    for k, v in state_dict.items():
        model_dict[k] = v.cpu().numpy()
    param_queue.put(model_dict)
    threading.Thread(target=_fill_param).start()
    learn_idx = 0
    ts = time.time()
    while True:
        #if batch_queue.empty():
        #    print("batch queue size:{}".format(batch_queue.qsize()))
        *batch, idxes = batch_queue.get()
        loss, prios = utils.compute_loss(model, tgt_model, batch, args.n_steps,
                                         args.gamma)
        grad_norm = utils.update_parameters(loss, model, optimizer,
                                            args.max_norm)
        prios_queue.put((idxes, prios))
        batch, idxes, prios = None, None, None
        learn_idx += 1

        if learn_idx % args.tensorboard_update_interval == 0:
            writer.add_scalar("learner/loss", loss, learn_idx)
            writer.add_scalar("learner/grad_norm", grad_norm, learn_idx)

        if learn_idx % args.target_update_interval == 0:
            print("Updating Target Network..")
            tgt_model.load_state_dict(model.state_dict())
        if learn_idx % args.save_interval == 0:
            print("Saving Model..")
            torch.save(model.state_dict(), "model.pth")
        if learn_idx % args.publish_param_interval == 0:
            param_queue.get()
        if learn_idx % args.bps_interval == 0:
            bps = args.bps_interval / (time.time() - ts)
            print("Step: {:8} / BPS: {:.2f}".format(learn_idx, bps))
            writer.add_scalar("learner/BPS", bps, learn_idx)
            ts = time.time()