def run()

in rlalgos/deprecated/sac/sac.py [0:0]

• 79 lines of code
• 6 McCabe index (conditional complexity)

    def run(self):
        self.replay_buffer = ReplayBuffer(self.config["replay_buffer_size"])
        device = torch.device(self.config["learner_device"])
        self.learning_model.to(device)

        self.q1.to(device)
        self.q2.to(device)
        self.target_q1.to(device)
        self.target_q2.to(device)
        optimizer = torch.optim.Adam(
            self.learning_model.parameters(), lr=self.config["lr"]
        )
        optimizer_q1 = torch.optim.Adam(self.q1.parameters(), lr=self.config["lr"])
        optimizer_q2 = torch.optim.Adam(self.q2.parameters(), lr=self.config["lr"])

        self.train_batcher.update(
            self._state_dict(self.learning_model, torch.device("cpu"))
        )
        self.evaluation_batcher.update(
            self._state_dict(self.learning_model, torch.device("cpu"))
        )

        n_episodes = self.config["n_envs"] * self.config["n_threads"]
        self.train_batcher.reset(
            agent_info=DictTensor({"stochastic": torch.zeros(n_episodes).eq(0.0)})
        )
        logging.info("Sampling initial transitions")
        n_iterations = int(
            self.config["n_starting_transitions"]
            / (n_episodes * self.config["batch_timesteps"])
        )
        for k in range(n_iterations):
            self.train_batcher.execute()
            trajectories = self.train_batcher.get()
            self.replay_buffer.push(trajectories)
        print("replay_buffer_size = ", self.replay_buffer.size())

        n_episodes = self.config["n_evaluation_rollouts"]
        stochastic = torch.tensor(
            [self.config["evaluation_mode"] == "stochastic"]
        ).repeat(n_episodes)
        self.evaluation_batcher.execute(
            agent_info=DictTensor({"stochastic": stochastic}), n_episodes=n_episodes
        )

        logging.info("Starting Learning")
        _start_time = time.time()

        logging.info("Learning")
        while time.time() - _start_time < self.config["time_limit"]:
            self.train_batcher.execute()
            trajectories = self.train_batcher.get()
            self.replay_buffer.push(trajectories)
            self.logger.add_scalar(
                "replay_buffer_size", self.replay_buffer.size(), self.iteration
            )
            # avg_reward = 0

            for k in range(self.config["n_batches_per_epochs"]):
                transitions = self.replay_buffer.sample(n=self.config["size_batches"])

                # print(dt)
                dt, transitions = self.get_q_loss(transitions, device)
                [
                    self.logger.add_scalar(k, dt[k].item(), self.iteration)
                    for k in dt.keys()
                ]
                optimizer_q1.zero_grad()
                dt["q1_loss"].backward()
                optimizer_q1.step()

                optimizer_q2.zero_grad()
                dt["q2_loss"].backward()
                optimizer_q2.step()

                optimizer.zero_grad()
                dt = self.get_policy_loss(transitions)
                [
                    self.logger.add_scalar(k, dt[k].item(), self.iteration)
                    for k in dt.keys()
                ]
                dt["policy_loss"].backward()
                optimizer.step()

                tau = self.config["tau"]
                self.soft_update_params(self.q1, self.target_q1, tau)
                self.soft_update_params(self.q2, self.target_q2, tau)

                self.iteration += 1

            self.train_batcher.update(
                self._state_dict(self.learning_model, torch.device("cpu"))
            )
            self.evaluate()