ray-project--ray/rllib/algorithms/iql/iql.py

from typing import Optional, Type, Union

from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig
from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
    AddObservationsFromEpisodesToBatch,
)
from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import (  # noqa
    AddNextObservationsFromEpisodesToTrainBatch,
)
from ray.rllib.core.learner.learner import Learner
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.utils.annotations import override
from ray.rllib.utils.typing import LearningRateOrSchedule, RLModuleSpecType


class IQLConfig(MARWILConfig):
    """Defines a configuration class from which a new IQL Algorithm can be built

    .. testcode::
        :skipif: True

        from ray.rllib.algorithms.iql import IQLConfig
        # Run this from the ray directory root.
        config = IQLConfig().training(actor_lr=0.00001, gamma=0.99)
        config = config.offline_data(
            input_="./rllib/offline/tests/data/pendulum/pendulum-v1_enormous")

        # Build an Algorithm object from the config and run 1 training iteration.
        algo = config.build()
        algo.train()

    .. testcode::
        :skipif: True

        from ray.rllib.algorithms.iql import IQLConfig
        from ray import tune
        config = IQLConfig()
        # Print out some default values.
        print(config.beta)
        # Update the config object.
        config.training(
            lr=tune.grid_search([0.001, 0.0001]), beta=0.75
        )
        # Set the config object's data path.
        # Run this from the ray directory root.
        config.offline_data(
            input_="./rllib/offline/tests/data/pendulum/pendulum-v1_enormous"
        )
        # Set the config object's env, used for evaluation.
        config.environment(env="Pendulum-v1")
        # Use to_dict() to get the old-style python config dict
        # when running with tune.
        tune.Tuner(
            "IQL",
            param_space=config.to_dict(),
        ).fit()
    """

    def __init__(self, algo_class=None):
        super().__init__(algo_class=algo_class or IQL)

        # fmt: off
        # __sphinx_doc_begin__
        # The temperature for the actor loss.
        self.beta = 0.1

        # The expectile to use in expectile regression.
        self.expectile = 0.8

        # The learning rates for the actor, critic and value network(s).
        self.actor_lr = 3e-4
        self.critic_lr = 3e-4
        self.value_lr = 3e-4
        # Set `lr` parameter to `None` and ensure it is not used.
        self.lr = None

        # If a twin-Q architecture should be used (advisable).
        self.twin_q = True

        # How often the target network should be updated.
        self.target_network_update_freq = 0
        # The weight for Polyak averaging.
        self.tau = 1.0

        # __sphinx_doc_end__
        # fmt: on

    @override(MARWILConfig)
    def training(
        self,
        *,
        twin_q: Optional[bool] = NotProvided,
        expectile: Optional[float] = NotProvided,
        actor_lr: Optional[LearningRateOrSchedule] = NotProvided,
        critic_lr: Optional[LearningRateOrSchedule] = NotProvided,
        value_lr: Optional[LearningRateOrSchedule] = NotProvided,
        target_network_update_freq: Optional[int] = NotProvided,
        tau: Optional[float] = NotProvided,
        **kwargs,
    ) -> "IQLConfig":
        """Sets the training related configuration.

        Args:
            beta: The temperature to scaling advantages in exponential terms.
                Must be >> 0.0. The higher this parameter the less greedy
                (exploitative) the policy becomes. It also means that the policy
                is fitting less to the best actions in the dataset.
            twin_q: If a twin-Q architecture should be used (advisable).
            expectile: The expectile to use in expectile regression for the value
                function. For high expectiles the value function tries to match
                the upper tail of the Q-value distribution.
            actor_lr: The learning rate for the actor network. Actor learning rates
                greater than critic learning rates work well in experiments.
            critic_lr: The learning rate for the Q-network. Critic learning rates
                greater than value function learning rates work well in experiments.
            value_lr: The learning rate for the value function network.
            target_network_update_freq: The number of timesteps in between the target
                Q-network is fixed. Note, too high values here could harm convergence.
                The target network is updated via Polyak-averaging.
            tau: The update parameter for Polyak-averaging of the target Q-network.
                The higher this value the faster the weights move towards the actual
                Q-network.

        Return:
            This updated `AlgorithmConfig` object.
        """
        super().training(**kwargs)

        if twin_q is not NotProvided:
            self.twin_q = twin_q
        if expectile is not NotProvided:
            self.expectile = expectile
        if actor_lr is not NotProvided:
            self.actor_lr = actor_lr
        if critic_lr is not NotProvided:
            self.critic_lr = critic_lr
        if value_lr is not NotProvided:
            self.value_lr = value_lr
        if target_network_update_freq is not NotProvided:
            self.target_network_update_freq = target_network_update_freq
        if tau is not NotProvided:
            self.tau = tau

        return self

    @override(MARWILConfig)
    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
        if self.framework_str == "torch":
            from ray.rllib.algorithms.iql.torch.iql_torch_learner import IQLTorchLearner

            return IQLTorchLearner
        else:
            raise ValueError(
                f"The framework {self.framework_str} is not supported. "
                "Use `'torch'` instead."
            )

    @override(MARWILConfig)
    def get_default_rl_module_spec(self) -> RLModuleSpecType:
        if self.framework_str == "torch":
            from ray.rllib.algorithms.iql.torch.default_iql_torch_rl_module import (
                DefaultIQLTorchRLModule,
            )

            return RLModuleSpec(module_class=DefaultIQLTorchRLModule)
        else:
            raise ValueError(
                f"The framework {self.framework_str} is not supported. "
                "Use `torch` instead."
            )

    @override(MARWILConfig)
    def build_learner_connector(
        self,
        input_observation_space,
        input_action_space,
        device=None,
    ):
        pipeline = super().build_learner_connector(
            input_observation_space=input_observation_space,
            input_action_space=input_action_space,
            device=device,
        )

        # Remove unneeded connectors from the MARWIL connector pipeline.
        pipeline.remove("AddOneTsToEpisodesAndTruncate")
        pipeline.remove("GeneralAdvantageEstimation")

        # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right
        # after the corresponding "add-OBS-..." default piece).
        pipeline.insert_after(
            AddObservationsFromEpisodesToBatch,
            AddNextObservationsFromEpisodesToTrainBatch(),
        )

        return pipeline

    @override(MARWILConfig)
    def validate(self) -> None:
        # Call super's validation method.
        super().validate()

        # Ensure hyperparameters are meaningful.
        if self.beta <= 0.0:
            self._value_error(
                "For meaningful results, `beta` (temperature) parameter must be >> 0.0!"
            )
        if not 0.0 < self.expectile < 1.0:
            self._value_error(
                "For meaningful results, `expectile` parameter must be in (0, 1)."
            )

    @property
    def _model_config_auto_includes(self):
        return super()._model_config_auto_includes | {"twin_q": self.twin_q}


class IQL(MARWIL):
    """Implicit Q-learning (derived from MARWIL).

    Uses MARWIL training step.
    """

    @classmethod
    @override(MARWIL)
    def get_default_config(cls) -> AlgorithmConfig:
        return IQLConfig()