Method

class Method(ABC):
    """Base class for all methods."""

    __version__ = "1.0.0"
    __author__ = "Nathan Gavenski"
    __method_name__ = "Abstract Method"

    def __init__(
        self,
        environment: Env,
        environment_parameters: Dict[str, Any],
        discrete_loss: nn.Module = nn.CrossEntropyLoss,
        continuous_loss: nn.Module = nn.MSELoss,
        optimizer_fn: optim.Optimizer = optim.Adam,
        activation: Callable[[torch.Tensor], torch.Tensor] = nn.LeakyReLU,
    ) -> None:
        """Initialize base class."""
        super().__init__()
        self.environment = environment
        self.discrete = isinstance(environment.action_space, spaces.Discrete)
        self.discrete |= isinstance(environment.action_space, gym_spaces.Discrete)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.visual = False

        self.observation_size = environment.observation_space.shape
        if len(self.observation_size) == 1:
            self.observation_size = self.observation_size[0]

        if self.discrete:
            self.action_size = environment.action_space.n
            self.loss_fn = discrete_loss()
        else:
            self.action_size = environment.action_space.shape[0]
            self.loss_fn = continuous_loss()

        policy = self.hyperparameters.get('policy', 'MlpPolicy')
        if policy == 'MlpPolicy':
            self.policy = MLP(
                self.observation_size, self.action_size,
                self.hyperparameters.get('hidden_dim', None),
                activation=activation
            )
        elif policy == 'MlpWithAttention':
            self.policy = MlpWithAttention(
                self.observation_size,
                self.action_size,
                self.hyperparameters.get('hidden_dim', None),
                activation=activation
            )
        elif policy == 'MlpAttention':
            self.policy = MlpAttention(
                self.observation_size, self.action_size,
                self.hyperparameters.get('hidden_dim', None),
                activation=activation
            )
        elif policy in ['CnnPolicy', 'ResnetPolicy', 'AttResnetPolicy']:
            self.visual = True
            if policy == 'CnnPolicy':
                encoder = CNN(self.observation_size)
            elif policy == 'ResnetPolicy':
                encoder = Resnet(self.observation_size)
            elif policy == 'AttResnetPolicy':
                encoder = ResnetWithAttention(self.observation_size)
            else:
                raise ValueError(f"Encoder {policy} not implemented, is it a typo?")

            with torch.no_grad():
                output = encoder(torch.zeros(1, *self.observation_size[::-1]))

            linear = []
            in_dim = output.shape[-1]
            for hidden_dim in self.hyperparameters.get("hidden_dim", []):
                linear.append(nn.Linear(in_dim, hidden_dim))
                linear.append(activation())
                linear.append(nn.Dropout(0.3))
                in_dim = hidden_dim
            linear.append(nn.Linear(in_dim, self.action_size))
            linear = nn.Sequential(*linear)
            self.policy = nn.Sequential(encoder, linear)

        self.optimizer_fn = optimizer_fn(
            self.policy.parameters(),
            lr=environment_parameters['lr']
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward method for the method.

        Args:
            x (torch.Tensor): input.

        Returns:
            x (torch.Tensor): logits output.
        """
        raise NotImplementedError()

    def predict(
        self,
        obs: Union[np.ndarray, torch.Tensor],
        transforms: Callable[[torch.Tensor], torch.Tensor] = None
    ) -> Union[List[Number], Number]:
        """Predict method.

        Args:
            obs (Union[np.ndarray, torch.Tensor]): input observation.
            transforms (Callable[torch.Tensor, torch.Tensor]): transform function
                to change data. If data is an image, the transforms parameter is
                required. Defaults to None.

        Returns:
            action (Union[List[Number], Number): predicted action.
        """
        self.policy.eval()

        if isinstance(obs, np.ndarray):
            if not self.visual:
                obs = torch.from_numpy(obs)
                if transforms is not None:
                    obs = transforms(obs)
                if len(obs.shape) == 1:
                    obs = obs[None]
            else:
                if transforms is None:
                    raise ValueError("Visual information requires transforms parameter.")
                obs = transforms(obs)
                if len(obs.shape) == 3:
                    obs = obs[None]

        obs = obs.to(self.device)

        with torch.no_grad():
            actions = self.forward(obs)
            actions = actions[0]

            if self.discrete:
                return torch.argmax(actions).cpu().numpy()
            return actions.cpu().numpy()

    def save(self, path: str = None) -> None:
        """Save all model weights.

        Args:
            path (str): where to save the models. Defaults to None.
        """
        raise NotImplementedError()

    def load(self, path: str = None) -> Self:
        """Load all model weights.

        Args:
            path (str): where to look for the model's weights. Defaults to None.
        """
        raise NotImplementedError()

    # pylint: disable=W0221
    def train(
        self,
        n_epochs: int,
        train_dataset: DataLoader,
        eval_dataset: DataLoader = None
    ) -> Self:
        """Train process.

        Args:
            n_epochs (int): amount of epoch to run.
            train_dataset (DataLoader): data to train.
            eval_dataset (DataLoader): data to eval. Defaults to None.

        Returns:
            method (Self): trained method.
        """
        raise NotImplementedError()

    def _train(self, dataset: DataLoader[BaselineDataset]) -> Metrics:
        """Train loop.

        Args:
            dataset (DataLoader): train data.
        """
        raise NotImplementedError()

    def _eval(self, dataset: DataLoader[BaselineDataset]) -> Metrics:
        """Evaluation loop.

        Args:
            dataset (DataLoader): data to eval.
        """
        raise NotImplementedError()

    def _enjoy(
        self,
        render: bool = False,
        teacher_reward: Number = None,
        random_reward: Number = None,
        gym_version: str = "newest",
        transforms: Callable[[torch.Tensor], torch.Tensor] = None
    ) -> Metrics:
        """Function for evaluation of the policy in the environment

        Args:
            render (bool): Whether it should render. Defaults to False.
            teacher_reward (Number): reward for teacher policy.
            random_reward (Number): reward for a random policy.
            gym_version (str): Which version of gym GymWrapper should use.
                Defaults to "newest" (gymnasium).
            transofrms (Callable[torch.Tensor, torch.Tensor]): torchvision
                functions to transform the data (only required for visual
                environments). Defaults to None.

        Returns:
            Metrics:
                aer (Number): average reward for 100 episodes.
                aer_std (Number): standard deviation for aer.
                performance (Number): if teacher_reward and random_reward are
                    informed than the performance metric is calculated.
                perforamance_std (Number): standard deviation for performance.
        """
        if self.hyperparameters.get("vectorized_enjoy", False):
            return self._vector_enjoy(
                teacher_reward=teacher_reward,
                random_reward=random_reward,
                gym_version=gym_version,
                transforms=transforms
            )
        return self._sequential_enjoy(
            render=render,
            teacher_reward=teacher_reward,
            random_reward=random_reward,
            gym_version=gym_version,
            transforms=transforms
        )


    def _sequential_enjoy(
        self,
        render: bool = False,
        teacher_reward: Number = None,
        random_reward: Number = None,
        gym_version: str = "newest",
        transforms: Callable[[torch.Tensor], torch.Tensor] = None
    ) -> Metrics:
        """Function for evaluation of the policy in the environment

        Args:
            render (bool): Whether it should render. Defaults to False.
            teacher_reward (Number): reward for teacher policy.
            random_reward (Number): reward for a random policy.
            gym_version (str): Which version of gym GymWrapper should use.
                Defaults to "newest" (gymnasium).
            transofrms (Callable[torch.Tensor, torch.Tensor]): torchvision
                functions to transform the data (only required for visual
                environments). Defaults to None.

        Returns:
            Metrics:
                aer (Number): average reward for 100 episodes.
                aer_std (Number): standard deviation for aer.
                performance (Number): if teacher_reward and random_reward are
                    informed than the performance metric is calculated.
                perforamance_std (Number): standard deviation for performance.
        """
        environment = GymWrapper(self.environment, gym_version)
        average_reward = []
        for _ in range(100):
            done = False
            obs = environment.reset()
            accumulated_reward = 0
            while not done:
                if render:
                    environment.render()
                action = self.predict(obs, transforms)
                obs, reward, done, *_ = environment.step(action)
                accumulated_reward += reward
            average_reward.append(accumulated_reward)

        metrics = average_episodic_reward(average_reward)
        if teacher_reward is not None and random_reward is not None:
            metrics.update(performance(average_reward, teacher_reward, random_reward))
        return metrics

    def _vector_enjoy(
        self,
        teacher_reward: Number = None,
        random_reward: Number = None,
        gym_version: str = "newest",
        transforms: Callable[[torch.Tensor], torch.Tensor] = None,
        n_envs: int = 10,
    ) -> Metrics:
        """Function for evaluation of the policy in the environment

        Args:
            render (bool): Whether it should render. Defaults to False.
            teacher_reward (Number): reward for teacher policy.
            random_reward (Number): reward for a random policy.
            gym_version (str): Which version of gym GymWrapper should use.
                Defaults to "newest" (gymnasium).
            transforms (Callable[torch.Tensor, torch.Tensor]): torchvision
                functions to transform the data (only required for visual
                environments). Defaults to None.
            n_envs (int): number of parallel environments. Defaults to 10.

        Returns:
            Metrics:
                aer (Number): average reward for 100 episodes.
                aer_std (Number): standard deviation for aer.
                performance (Number): if teacher_reward and random_reward are
                    informed than the performance metric is calculated.
                performance_std (Number): standard deviation for performance.
        """
        from gymnasium.vector import SyncVectorEnv

        n_episodes = 100
        make_env = lambda: gymnasium.make(self.environment.spec.id)
        vec_env = SyncVectorEnv([make_env for _ in range(n_envs)])

        average_reward = []
        accumulated_rewards = np.zeros(n_envs)
        episodes_done = 0

        obs, _ = vec_env.reset()
        while episodes_done < n_episodes:
            actions = np.array([self.predict(o, transforms) for o in obs])
            obs, rewards, terminated, truncated, _ = vec_env.step(actions)
            accumulated_rewards += rewards

            done = terminated | truncated
            for i, d in enumerate(done):
                if d and episodes_done < n_episodes:
                    average_reward.append(accumulated_rewards[i])
                    accumulated_rewards[i] = 0
                    episodes_done += 1

        vec_env.close()

        metrics = average_episodic_reward(average_reward)
        if teacher_reward is not None and random_reward is not None:
            metrics.update(performance(average_reward, teacher_reward, random_reward))
        return metrics

    def early_stop(self, metric: Metrics) -> bool:
        """Function that tells the method if it should stop or not.

        Args:
            metric (Metrics): the metric to evaluate whether it should early stop.

        Returns:
            stop (bool): if it should stop or not.
        """
        return False