Ppo

mighty.mighty_agents.ppo

MightyPPOAgent

MightyPPOAgent(
    output_dir,
    env: MIGHTYENV,
    eval_env: Optional[MIGHTYENV] = None,
    seed: Optional[int] = None,
    learning_rate: float = 0.001,
    gamma: float = 0.99,
    batch_size: int = 64,
    learning_starts: int = 1,
    render_progress: bool = True,
    log_wandb: bool = False,
    wandb_kwargs: Optional[Dict] = {
        "project": "mighty",
        "entity": "amsks",
        "name": "ppo",
        "group": "ppo",
    },
    rollout_buffer_class: Optional[
        str | DictConfig | Type[MightyRolloutBuffer]
    ] = MightyRolloutBuffer,
    rollout_buffer_kwargs: Optional[TypeKwargs] = None,
    meta_methods: Optional[List[str | type]] = None,
    meta_kwargs: Optional[List[TypeKwargs]] = None,
    n_policy_units: int = 8,
    n_critic_units: int = 8,
    soft_update_weight: float = 0.01,
    policy_class: Optional[
        Union[
            str, DictConfig, Type[MightyExplorationPolicy]
        ]
    ] = None,
    policy_kwargs: Optional[Dict] = None,
    ppo_clip: float = 0.2,
    value_loss_coef: float = 0.5,
    entropy_coef: float = 0.01,
    max_grad_norm: float = 0.5,
    n_gradient_steps: int = 10,
    hidden_sizes: Optional[List[int]] = [64, 64],
    activation: Optional[str] = "tanh",
    n_epochs: int = 10,
    minibatch_size: int = 32,
    kl_target: float = 0.001,
    use_value_clip: bool = True,
    value_clip_eps: float = 0.2,
    total_timesteps: int = 1000000,
)

Bases: MightyAgent

Creates all relevant class variables and calls the agent-specific init function.

:param env: Train environment :param eval_env: Evaluation environment :param seed: Seed for random number generators :param learning_rate: Learning rate for training :param gamma: Discount factor :param batch_size: Batch size for training :param learning_starts: Number of steps before learning starts :param render_progress: Whether to render progress :param log_tensorboard: Log to TensorBoard as well as to file :param log_wandb: Log to Weights and Biases :param wandb_kwargs: Arguments for Weights and Biases logging :param rollout_buffer_class: Rollout buffer class :param rollout_buffer_kwargs: Arguments for the rollout buffer :param meta_methods: Meta methods for the agent :param meta_kwargs: Arguments for meta methods :param n_policy_units: Number of units for the policy network :param n_critic_units: Number of units for the critic network :param soft_update_weight: Size of soft updates for the target network :param policy_class: Policy class :param policy_kwargs: Arguments for the policy :param ppo_clip: Clipping parameter for PPO :param value_loss_coef: Coefficient for the value loss :param entropy_coef: Coefficient for the entropy loss :param max_grad_norm: Maximum gradient norm :param n_gradient_steps: Number of gradient steps per update

Source code in mighty/mighty_agents/ppo.py

def __init__(
    self,
    output_dir,
    env: MIGHTYENV,  # type: ignore
    eval_env: Optional[MIGHTYENV] = None,  # type: ignore
    seed: Optional[int] = None,
    learning_rate: float = 0.001,
    gamma: float = 0.99,
    batch_size: int = 64,
    learning_starts: int = 1,
    render_progress: bool = True,
    log_wandb: bool = False,
    wandb_kwargs: Optional[Dict] = {
        "project": "mighty",
        "entity": "amsks",
        "name": "ppo",
        "group": "ppo",
    },
    rollout_buffer_class: Optional[
        str | DictConfig | Type[MightyRolloutBuffer]
    ] = MightyRolloutBuffer,
    rollout_buffer_kwargs: Optional[TypeKwargs] = None,
    meta_methods: Optional[List[str | type]] = None,
    meta_kwargs: Optional[List[TypeKwargs]] = None,
    n_policy_units: int = 8,
    n_critic_units: int = 8,
    soft_update_weight: float = 0.01,
    policy_class: Optional[
        Union[str, DictConfig, Type[MightyExplorationPolicy]]
    ] = None,
    policy_kwargs: Optional[Dict] = None,
    ppo_clip: float = 0.2,
    value_loss_coef: float = 0.5,
    entropy_coef: float = 0.01,
    max_grad_norm: float = 0.5,
    n_gradient_steps: int = 10,
    hidden_sizes: Optional[List[int]] = [64, 64],
    activation: Optional[str] = "tanh",
    n_epochs: int = 10,
    minibatch_size: int = 32,
    kl_target: float = 0.001,
    use_value_clip: bool = True,
    value_clip_eps: float = 0.2,
    total_timesteps: int = 1_000_000,
):
    """Initialize the PPO agent.

    Creates all relevant class variables and calls the agent-specific init function.

    :param env: Train environment
    :param eval_env: Evaluation environment
    :param seed: Seed for random number generators
    :param learning_rate: Learning rate for training
    :param gamma: Discount factor
    :param batch_size: Batch size for training
    :param learning_starts: Number of steps before learning starts
    :param render_progress: Whether to render progress
    :param log_tensorboard: Log to TensorBoard as well as to file
    :param log_wandb: Log to Weights and Biases
    :param wandb_kwargs: Arguments for Weights and Biases logging
    :param rollout_buffer_class: Rollout buffer class
    :param rollout_buffer_kwargs: Arguments for the rollout buffer
    :param meta_methods: Meta methods for the agent
    :param meta_kwargs: Arguments for meta methods
    :param n_policy_units: Number of units for the policy network
    :param n_critic_units: Number of units for the critic network
    :param soft_update_weight: Size of soft updates for the target network
    :param policy_class: Policy class
    :param policy_kwargs: Arguments for the policy
    :param ppo_clip: Clipping parameter for PPO
    :param value_loss_coef: Coefficient for the value loss
    :param entropy_coef: Coefficient for the entropy loss
    :param max_grad_norm: Maximum gradient norm
    :param n_gradient_steps: Number of gradient steps per update
    """

    self.total_timesteps = total_timesteps
    self.gamma = gamma
    self.n_policy_units = n_policy_units
    self.n_critic_units = n_critic_units
    self.soft_update_weight = soft_update_weight
    self.ppo_clip = ppo_clip
    self.value_loss_coef = value_loss_coef
    self.entropy_coef = entropy_coef
    self.max_grad_norm = max_grad_norm
    self.hidden_sizes = hidden_sizes
    self.activation = activation

    self.n_epochs = n_epochs
    self.minibatch_size = minibatch_size
    self.kl_target = kl_target
    self.use_value_clip = use_value_clip
    self.value_clip_eps = value_clip_eps

    # Placeholder variables which are filled in self._initialize_agent
    self.model: PPOModel | None = None
    self.update_fn: PPOUpdate | None = None

    # Policy class
    policy_class = retrieve_class(cls=policy_class, default_cls=StochasticPolicy)  # type: ignore
    if policy_kwargs is None:
        policy_kwargs = {}
    self.policy_class = policy_class
    self.policy_kwargs = policy_kwargs

    super().__init__(
        env=env,
        output_dir=output_dir,
        seed=seed,
        eval_env=eval_env,
        learning_rate=learning_rate,
        batch_size=batch_size,
        learning_starts=learning_starts,
        n_gradient_steps=n_gradient_steps,
        render_progress=render_progress,
        log_wandb=log_wandb,
        wandb_kwargs=wandb_kwargs,
        replay_buffer_class=rollout_buffer_class,
        replay_buffer_kwargs=rollout_buffer_kwargs,
        meta_methods=meta_methods,
        meta_kwargs=meta_kwargs,
    )

    self.loss_buffer = {
        "Update/policy_loss": [],
        "Update/value_loss": [],
        "Update/entropy": [],
        "Update/approx_kl": [],
        "step": [],
    }

    if self.log_wandb:
        wandb.init(**(wandb_kwargs or {}))

agent_type property

agent_type: str

Return the type of the agent.

value_function property

value_function: Module

Return the value function model.

__del__

__del__() -> None

Close wandb upon deletion.

Source code in mighty/mighty_agents/base_agent.py

def __del__(self) -> None:
    """Close wandb upon deletion."""
    self.env.close()  # type: ignore
    if self.log_wandb:
        wandb.finish()

adapt_hps

adapt_hps(metrics: Dict) -> None

Set hyperparameters.

Source code in mighty/mighty_agents/base_agent.py

def adapt_hps(self, metrics: Dict) -> None:
    """Set hyperparameters."""
    old_hps = {
        "step": self.steps,
        "hp/lr": self.learning_rate,
        "hp/pi_epsilon": self._epsilon,
        "hp/batch_size": self._batch_size,
        "hp/learning_starts": self._learning_starts,
        "meta_modules": list(self.meta_modules.keys()),
    }
    self.learning_rate = metrics["hp/lr"]
    self._epsilon = metrics["hp/pi_epsilon"]
    self._batch_size = metrics["hp/batch_size"]
    self._learning_starts = metrics["hp/learning_starts"]

    updated_hps = {
        "step": self.steps,
        "hp/lr": self.learning_rate,
        "hp/pi_epsilon": self._epsilon,
        "hp/batch_size": self._batch_size,
        "hp/learning_starts": self._learning_starts,
        "meta_modules": list(self.meta_modules.keys()),
    }
    # FIXME: EWRL: Can we naively compare dictoinaries like this?
    if old_hps != updated_hps:
        self.hp_buffer = update_buffer(self.hp_buffer, updated_hps)

apply_config

apply_config(config: Dict) -> None

Apply config to agent.

Source code in mighty/mighty_agents/base_agent.py

def apply_config(self, config: Dict) -> None:
    """Apply config to agent."""
    for n in config:
        algo_name = n.split(".")[-1]
        if hasattr(self, algo_name):
            setattr(self, algo_name, config[n])
        elif hasattr(self, "_" + algo_name):
            setattr(self, "_" + algo_name, config[n])
        elif n in ["architecture", "n_units", "n_layers", "size"]:
            pass
        else:
            print(f"Trying to set hyperparameter {algo_name} which does not exist.")

evaluate

evaluate(eval_env: MIGHTYENV | None = None) -> Dict

Eval agent on an environment. (Full rollouts).

:param env: The environment to evaluate on :param episodes: The number of episodes to evaluate :return:

Source code in mighty/mighty_agents/base_agent.py

def evaluate(self, eval_env: MIGHTYENV | None = None) -> Dict:  # type: ignore
    """Eval agent on an environment. (Full rollouts).

    :param env: The environment to evaluate on
    :param episodes: The number of episodes to evaluate
    :return:
    """

    terminated, truncated = False, False
    options: Dict = {}
    if eval_env is None:
        eval_env = self.eval_env

    state, _ = eval_env.reset(options=options)  # type: ignore
    rewards = np.zeros(eval_env.num_envs)  # type: ignore
    steps = np.zeros(eval_env.num_envs)  # type: ignore
    mask = np.zeros(eval_env.num_envs)  # type: ignore
    while not np.all(mask):
        action = self.policy(state, evaluate=True)  # type: ignore
        state, reward, terminated, truncated, _ = eval_env.step(action)  # type: ignore
        rewards += reward * (1 - mask)
        steps += 1 * (1 - mask)
        dones = np.logical_or(terminated, truncated)
        mask = np.where(dones, 1, mask)

    eval_env.close()  # type: ignore

    if isinstance(self.eval_env, DACENV) or isinstance(self.env, CARLENV):
        instance = eval_env.instance  # type: ignore
    else:
        instance = "None"

    eval_metrics = {
        "step": self.steps,
        "seed": self.seed,
        "eval_episodes": np.array(rewards) / steps,
        "mean_eval_step_reward": np.mean(rewards) / steps,
        "mean_eval_reward": np.mean(rewards),
        "instance": instance,
        "eval_rewards": rewards,
    }
    self.eval_buffer = update_buffer(self.eval_buffer, eval_metrics)

    if self.log_wandb:
        wandb.log(eval_metrics)

    return eval_metrics

initialize_agent

initialize_agent() -> None

General initialization of tracer and buffer for all agents.

Algorithm specific initialization like policies etc. are done in _initialize_agent

Source code in mighty/mighty_agents/base_agent.py

def initialize_agent(self) -> None:
    """General initialization of tracer and buffer for all agents.

    Algorithm specific initialization like policies etc.
    are done in _initialize_agent
    """
    self._initialize_agent()
    self.buffer = self.buffer_class(**self.buffer_kwargs)  # type: ignore

load

load(path: str) -> None

Load the internal state of the agent.

Source code in mighty/mighty_agents/ppo.py

def load(self, path: str) -> None:
    """Load the internal state of the agent."""
    base_path = Path(path)
    self.model.policy_head.load_state_dict(torch.load(base_path / "policy_head.pt"))  # type: ignore
    self.model.value_head.load_state_dict(torch.load(base_path / "value_head.pt"))  # type: ignore
    self.update_fn.policy_optimizer.load_state_dict(  # type: ignore
        torch.load(base_path / "policy_optimizer.pt")
    )
    self.update_fn.value_optimizer.load_state_dict(  # type: ignore
        torch.load(base_path / "value_optimizer.pt")
    )

    if self.verbose:
        print(f"Loaded checkpoint at {path}")

make_checkpoint_dir

make_checkpoint_dir(t: int) -> None

Checkpoint model.

:param T: Current timestep :return:

Source code in mighty/mighty_agents/base_agent.py

def make_checkpoint_dir(self, t: int) -> None:
    """Checkpoint model.

    :param T: Current timestep
    :return:
    """
    self.upper_checkpoint_dir = Path(self.output_dir) / Path("checkpoints")
    if not self.upper_checkpoint_dir.exists():
        Path(self.upper_checkpoint_dir).mkdir()
    self.checkpoint_dir = self.upper_checkpoint_dir / f"{t}"
    if not self.checkpoint_dir.exists():
        Path(self.checkpoint_dir).mkdir()

run

run(
    n_steps: int,
    eval_every_n_steps: int = 1000,
    save_model_every_n_steps: int | None = 5000,
    env: MIGHTYENV = None,
) -> Dict

Run agent.

Source code in mighty/mighty_agents/base_agent.py

def run(  # noqa: PLR0915
    self,
    n_steps: int,
    eval_every_n_steps: int = 1_000,
    save_model_every_n_steps: int | None = 5000,
    env: MIGHTYENV = None,  # type: ignore
) -> Dict:
    """Run agent."""
    episodes = 0
    if env is not None:
        self.env = env

    logging_layout = Layout()
    logging_layout.split_column(
        Layout(name="upper"), Layout(name="middle"), Layout(name="lower")
    )
    logging_layout["middle"].split_row(
        Layout(name="left"),
        Layout(name="right"),
    )
    logging_layout["lower"].split_row(
        Layout(name="left"),
        Layout(name="right"),
    )

    metrics_table = self.make_logging_table(0, [0], [0], [0], [0])  # type: ignore
    training_detail_table = Table(title="Training Details")
    training_detail_table.add_column("Name", style="cyan", justify="center")
    training_detail_table.add_column("Value", style="cyan", justify="center")
    training_detail_table.add_row("Training Steps", str(n_steps))
    training_detail_table.add_row("Seed", str(self.seed))
    training_detail_table.add_row("Learning Rate", str(self.learning_rate))
    training_detail_table.add_row("Batch Size", str(self._batch_size))

    logging_layout["middle"]["left"].update(metrics_table)
    logging_layout["middle"]["right"].update(training_detail_table)

    progress = Progress(
        "[progress.description]{task.description}",
        BarColumn(),
        "[progress.percentage]{task.percentage:>3.0f}%",
        "Remaining:",
        TimeRemainingColumn(),
        "Elapsed:",
        TimeElapsedColumn(),
        disable=not self.render_progress,
    )
    steps_task = progress.add_task(
        "Train Steps",
        total=n_steps - self.steps,
        start=False,
        visible=False,
    )
    progress.start_task(steps_task)

    progress_table = Table.grid(expand=True)
    progress_table.add_row(
        Panel(
            progress,
            title="Training Progress",
            border_style="green",
            padding=(2, 2),
        ),
    )
    logging_layout["upper"].update(progress_table)
    update_multiplier = 0

    with Live(logging_layout, refresh_per_second=10, vertical_overflow="visible"):
        steps_since_eval = 0
        # FIXME: this is more of a question: are there cases where we don't want to reset this completely?
        # I can't think of any, can you? If yes, we should maybe add this as an optional argument
        metrics = {
            "env": self.env,
            "vf": self.value_function,  # type: ignore
            "policy": self.policy,
            "step": self.steps,
            "hp/lr": self.learning_rate,
            "hp/pi_epsilon": self._epsilon,
            "hp/batch_size": self._batch_size,
            "hp/learning_starts": self._learning_starts,
        }

        # Reset env and initialize reward sum
        curr_s, _ = self.env.reset()  # type: ignore
        if len(curr_s.squeeze().shape) == 0:
            episode_reward = [0]
        else:
            episode_reward = np.zeros(curr_s.squeeze().shape[0])  # type: ignore

        last_episode_reward = episode_reward
        if not torch.is_tensor(last_episode_reward):
            last_episode_reward = torch.tensor(last_episode_reward).float()
        progress.update(steps_task, visible=True)
        recent_episode_reward = []
        recent_step_reward = []
        recent_actions = []
        evaluation_reward = []
        eval_curve = [0]
        learning_curve = [0]
        curve_xs = [0]
        logging_layout["lower"]["left"].update(
            self.get_plot(curve_xs, learning_curve, "Training Reward")
        )
        logging_layout["lower"]["right"].update(
            self.get_plot(curve_xs, eval_curve, "Evaluation Reward")
        )

        # Main loop: rollouts, training and evaluation
        while self.steps < n_steps:
            metrics["episode_reward"] = episode_reward

            action, log_prob = self.step(curr_s, metrics)

            next_s, reward, terminated, truncated, _ = self.env.step(action)  # type: ignore
            dones = np.logical_or(terminated, truncated)

            episode_reward += reward

            # Log everything
            t = {
                "seed": self.seed,
                "step": self.steps,
                "reward": reward,
                "action": action,
                "state": curr_s,
                "next_state": next_s,
                "terminated": terminated.astype(int),
                "truncated": truncated.astype(int),
                "dones": dones.astype(int),
                "mean_episode_reward": last_episode_reward.mean().cpu().numpy().item(),
            }
            metrics["log_prob"] = log_prob.detach().cpu().numpy()
            metrics["episode_reward"] = episode_reward
            metrics["transition"] = t

            recent_actions.append(np.mean(action))
            if len(recent_actions) > 100:
                recent_actions.pop(0)

            for k in self.meta_modules:
                self.meta_modules[k].post_step(metrics)

            transition_metrics = self.process_transition(
                metrics["transition"]["state"],
                metrics["transition"]["action"],
                metrics["transition"]["reward"],
                metrics["transition"]["next_state"],
                metrics["transition"]["dones"],
                metrics["log_prob"],
                metrics,
            )
            metrics.update(transition_metrics)

            self.result_buffer = update_buffer(self.result_buffer, t)

            if self.log_wandb:
                wandb.log(t)

            self.steps += len(action)
            metrics["step"] = self.steps
            steps_since_eval += len(action)
            for _ in range(len(action)):
                progress.advance(steps_task)

            # Update agent
            if (
                len(self.buffer) >= self._batch_size  # type: ignore
                and self.steps >= self._learning_starts
            ):
                update_kwargs = {"next_s": next_s, "dones": dones}

                metrics = self.update(metrics, update_kwargs)

            # End step
            self.last_state = curr_s
            curr_s = next_s

            # Evaluate
            if eval_every_n_steps and steps_since_eval >= eval_every_n_steps:
                steps_since_eval = 0
                eval_metrics = self.evaluate()
                evaluation_reward = eval_metrics["eval_rewards"]

            # Log to command line
            if self.steps >= 1000*update_multiplier:
                metrics_table = self.make_logging_table(
                    self.steps,
                    recent_episode_reward,
                    recent_step_reward,
                    evaluation_reward,
                    recent_actions,
                )
                logging_layout["middle"]["left"].update(metrics_table)
                eval_curve.append(np.mean(evaluation_reward))
                learning_curve.append(np.mean(recent_episode_reward))
                curve_xs.append(self.steps)

                logging_layout["lower"]["left"].update(
                    self.get_plot(curve_xs, learning_curve, "Training Reward")
                )
                logging_layout["lower"]["right"].update(
                    self.get_plot(curve_xs, eval_curve, "Evaluation Reward")
                )
                update_multiplier += 1

            # Save
            if (
                save_model_every_n_steps
                and self.steps % save_model_every_n_steps == 0
            ):
                self.save(self.steps)
                log_to_file(
                    self.output_dir,
                    self.result_buffer,
                    self.hp_buffer,
                    self.eval_buffer,
                    self.loss_buffer,
                )

            if np.any(dones):
                last_episode_reward = np.where(  # type: ignore
                    dones, episode_reward, last_episode_reward
                )
                recent_episode_reward.append(np.mean(last_episode_reward))
                recent_step_reward.append(
                    np.mean(last_episode_reward) / len(last_episode_reward)
                )
                last_episode_reward = torch.tensor(last_episode_reward).float()
                if len(recent_episode_reward) > 10:
                    recent_episode_reward.pop(0)
                    recent_step_reward.pop(0)
                episode_reward = np.where(dones, 0, episode_reward)  # type: ignore
                # End episode
                if isinstance(self.env, DACENV) or isinstance(self.env, CARLENV):
                    instance = self.env.instance  # type: ignore
                else:
                    instance = None
                metrics["instance"] = instance
                episodes += 1
                for k in self.meta_modules:
                    self.meta_modules[k].post_episode(metrics)

                # Meta Module hooks
                for k in self.meta_modules:
                    self.meta_modules[k].pre_episode(metrics)
    log_to_file(
        self.output_dir,
        self.result_buffer,
        self.hp_buffer,
        self.eval_buffer,
        self.loss_buffer,
    )
    return metrics

save

save(t: int) -> None

Save current agent state.

Source code in mighty/mighty_agents/ppo.py

def save(self, t: int) -> None:
    """Save current agent state."""
    super().make_checkpoint_dir(t)
    torch.save(
        self.model.policy_head.state_dict(),  # type: ignore
        self.checkpoint_dir / "policy_head.pt",
    )
    torch.save(
        self.model.value_head.state_dict(),  # type: ignore
        self.checkpoint_dir / "value_head.pt",
    )
    torch.save(
        self.update_fn.optimizer.state_dict(),  # type: ignore
        self.checkpoint_dir / "optimizer.pt",
    )

    if self.verbose:
        print(f"Saved checkpoint at {self.checkpoint_dir}")

update_agent

update_agent(
    transition_batch, batches_left, next_s, dones, **kwargs
) -> Dict

Update the agent using PPO.

:return: Dictionary containing the update metrics.

Source code in mighty/mighty_agents/ppo.py

def update_agent(
    self, transition_batch, batches_left, next_s, dones, **kwargs
) -> Dict:  # type: ignore
    """Update the agent using PPO.

    :return: Dictionary containing the update metrics.
    """
    metrics: Dict = {}
    metrics.update(self.update_fn.update(transition_batch))  # type: ignore

    for key, value in metrics.items():
        self.loss_buffer[key].append(value)
    self.loss_buffer["step"].append(self.steps)

    # Wandb logging
    if self.log_wandb:
        serializable_metrics = {}
        for k, v in metrics.items():
            try:
                json.dumps(v)
                serializable_metrics[k] = v
            except TypeError:
                print(f"Skipping non-serializable metric: {k}")

        wandb.log(serializable_metrics, step=self.steps)

    if batches_left == 0:
        self.buffer.reset()  # type: ignore

    return metrics