chore: import upstream snapshot with attribution

rllib/algorithms/iql/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.algorithms.iql.iql import IQL, IQLConfig
+
+__all__ = [
+    "IQL",
+    "IQLConfig",
+]

rllib/algorithms/iql/default_iql_rl_module.py

@@ -0,0 +1,35 @@
+from ray.rllib.algorithms.sac.default_sac_rl_module import DefaultSACRLModule
+from ray.rllib.core.models.configs import MLPHeadConfig
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+
+
+class DefaultIQLRLModule(DefaultSACRLModule, ValueFunctionAPI):
+    @override(DefaultSACRLModule)
+    def setup(self):
+        # Setup the `DefaultSACRLModule` to get the catalog.
+        super().setup()
+
+        # Only, if the `RLModule` is used on a `Learner` we build the value network.
+        if not self.inference_only:
+            # Build the encoder for the value function.
+            self.vf_encoder = self.catalog.build_encoder(framework=self.framework)
+
+            # Build the vf head.
+            self.vf = MLPHeadConfig(
+                input_dims=self.catalog.latent_dims,
+                # Note, we use the same layers as for the policy and Q-network.
+                hidden_layer_dims=self.catalog.pi_and_qf_head_hiddens,
+                hidden_layer_activation=self.catalog.pi_and_qf_head_activation,
+                output_layer_activation="linear",
+                output_layer_dim=1,
+            ).build(framework=self.framework)
+
+    @override(DefaultSACRLModule)
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    def get_non_inference_attributes(self):
+        # Use all of `super`'s attributes and add the value function attributes.
+        return super().get_non_inference_attributes() + ["vf_encoder", "vf"]

rllib/algorithms/iql/iql.py

@@ -0,0 +1,228 @@
+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()

rllib/algorithms/iql/iql_learner.py

@@ -0,0 +1,84 @@
+from typing import Dict
+
+from ray.rllib.algorithms.dqn.dqn_learner import DQNLearner
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+QF_TARGET_PREDS = "qf_target_preds"
+VF_PREDS_NEXT = "vf_preds_next"
+VF_LOSS = "value_loss"
+
+
+class IQLLearner(DQNLearner):
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(DQNLearner)
+    def build(self) -> None:
+        # Build the `DQNLearner` (builds the target network).
+        super().build()
+
+        # Define the expectile parameter(s).
+        self.expectile: Dict[ModuleID, TensorType] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                # Note, we want to train with a certain expectile.
+                [self.config.get_config_for_module(module_id).expectile],
+                trainable=False,
+            )
+        )
+
+        # Define the temperature for the actor advantage loss.
+        self.temperature: Dict[ModuleID, TensorType] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                # Note, we want to train with a certain expectile.
+                [self.config.get_config_for_module(module_id).beta],
+                trainable=False,
+            )
+        )
+
+        # Store loss tensors here temporarily inside the loss function for (exact)
+        # consumption later by the compute gradients function.
+        # Keys=(module_id, optimizer_name), values=loss tensors (in-graph).
+        self._temp_losses = {}
+
+    @override(DQNLearner)
+    def remove_module(self, module_id: ModuleID) -> None:
+        """Removes the expectile and temperature for removed modules."""
+        # First call `super`'s `remove_module` method.
+        super().remove_module(module_id)
+        # Remove the expectile from the mapping.
+        self.expectile.pop(module_id, None)
+        # Remove the temperature from the mapping.
+        self.temperature.pop(module_id, None)
+
+    @override(DQNLearner)
+    def add_module(
+        self,
+        *,
+        module_id,
+        module_spec,
+        config_overrides=None,
+        new_should_module_be_updated=None
+    ):
+        """Adds the expectile and temperature for new modules."""
+        # First call `super`'s `add_module` method.
+        super().add_module(
+            module_id=module_id,
+            module_spec=module_spec,
+            config_overrides=config_overrides,
+            new_should_module_be_updated=new_should_module_be_updated,
+        )
+        # Add the expectile to the mapping.
+        self.expectile[module_id] = self._get_tensor_variable(
+            # Note, we want to train with a certain expectile.
+            [self.config.get_config_for_module(module_id).beta],
+            trainable=False,
+        )
+        # Add the temperature to the mapping.
+        self.temperature[module_id] = self._get_tensor_variable(
+            # Note, we want to train with a certain expectile.
+            [self.config.get_config_for_module(module_id).beta],
+            trainable=False,
+        )

rllib/algorithms/iql/torch/default_iql_torch_rl_module.py

@@ -0,0 +1,79 @@
+from typing import Any, Dict, Optional
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.iql.default_iql_rl_module import DefaultIQLRLModule
+from ray.rllib.algorithms.iql.iql_learner import QF_TARGET_PREDS, VF_PREDS_NEXT
+from ray.rllib.algorithms.sac.torch.default_sac_torch_rl_module import (
+    DefaultSACTorchRLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class DefaultIQLTorchRLModule(DefaultSACTorchRLModule, DefaultIQLRLModule):
+
+    framework: str = "torch"
+
+    @override(DefaultSACTorchRLModule)
+    def _forward_train(self, batch: Dict, **kwargs) -> Dict[str, Any]:
+
+        # Right now, IQL runs only with continuous action spaces.
+        # TODO (simon): Implement it also for discrete action spaces.
+        if not isinstance(self.action_space, gym.spaces.Box):
+            raise ValueError(
+                f"Unsupported action space type: {type(self.action_space)}. "
+                "Only continuous action spaces are supported."
+            )
+
+        # Call the forward pass of the SAC module.
+        output = super()._forward_train(batch, **kwargs)
+
+        # Create batches for the forward passes of the target Q-networks and the
+        # value function.
+        batch_curr = {
+            Columns.OBS: batch[Columns.OBS],
+            Columns.ACTIONS: batch[Columns.ACTIONS],
+        }
+        batch_next = {Columns.OBS: batch[Columns.NEXT_OBS]}
+
+        # These target q-values are needed for the value loss and actor loss.
+        output[QF_TARGET_PREDS] = self._qf_forward_train_helper(
+            batch_curr, encoder=self.target_qf_encoder, head=self.target_qf
+        )
+        # If a twin-Q architecture is used run its target Q-network.
+        if self.twin_q:
+            output[QF_TARGET_PREDS] = torch.min(
+                output[QF_TARGET_PREDS],
+                self._qf_forward_train_helper(
+                    batch_curr, encoder=self.target_qf_twin_encoder, head=self.qf_twin
+                ),
+            )
+
+        # Compute values for the current observations.
+        output[Columns.VF_PREDS] = self.compute_values(batch_curr)
+        # The values of the next observations are needed for the critic loss.
+        output[VF_PREDS_NEXT] = self.compute_values(batch_next)
+
+        return output
+
+    @override(ValueFunctionAPI)
+    def compute_values(
+        self,
+        batch: Dict[str, Any],
+        embeddings: Optional[Any] = None,
+    ) -> TensorType:
+        # If no embeddings are provided make a forward pass on the encoder.
+        if embeddings is None:
+            embeddings = self.vf_encoder(batch)[ENCODER_OUT]
+
+        # Value head.
+        vf_out = self.vf(embeddings)
+        # Squeeze out last dimension (single node value head).
+        return vf_out.squeeze(-1)

rllib/algorithms/iql/torch/iql_torch_learner.py

@@ -0,0 +1,245 @@
+from typing import Dict
+
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.dqn.dqn_learner import QF_LOSS_KEY, QF_PREDS
+from ray.rllib.algorithms.iql.iql_learner import (
+    QF_TARGET_PREDS,
+    VF_LOSS,
+    VF_PREDS_NEXT,
+    IQLLearner,
+)
+from ray.rllib.algorithms.sac.sac_learner import QF_TWIN_LOSS_KEY, QF_TWIN_PREDS
+from ray.rllib.core import ALL_MODULES
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import (
+    POLICY_LOSS_KEY,
+)
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID, ParamDict, TensorType
+
+torch, nn = try_import_torch()
+
+
+class IQLTorchLearner(TorchLearner, IQLLearner):
+    """Implements the IQL loss on top of `IQLLearner`.
+
+    This Learner implements configure_optimizers_for_module to define
+    separate optimizers for the policy, Q-, and value networks. When
+    using a twin-Q network architecture, each Q-network is assigned its
+    own optimizer—consistent with the SAC algorithm.
+
+    The IQL loss is defined in compute_loss_for_module and consists of
+    three components: value loss, Q-loss (TD error), and actor (policy)
+    loss.
+
+    Note that the original IQL implementation performs separate backward
+    passes for each network. However, due to RLlib's reliance on TorchDDP,
+    all backward passes must be executed within a single update step. This
+    constraint can lead to parameter lag and cyclical loss behavior, though
+    it does not hinder convergence.
+    """
+
+    @override(TorchLearner)
+    def configure_optimizers_for_module(
+        self, module_id: ModuleID, config: AlgorithmConfig = None
+    ) -> None:
+
+        # Note, we could have derived directly from SACTorchLearner to
+        # inherit the setup of optimizers, but that learner comes with
+        # additional parameters which we do not need.
+        # Receive the module.
+        module = self._module[module_id]
+
+        # Define the optimizer for the critic.
+        # TODO (sven): Maybe we change here naming to `qf` for unification.
+        params_critic = self.get_parameters(module.qf_encoder) + self.get_parameters(
+            module.qf
+        )
+        optim_critic = torch.optim.Adam(params_critic, eps=1e-7)
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="qf",
+            optimizer=optim_critic,
+            params=params_critic,
+            lr_or_lr_schedule=config.critic_lr,
+        )
+        # If necessary register also an optimizer for a twin Q network.
+        if config.twin_q:
+            params_twin_critic = self.get_parameters(
+                module.qf_twin_encoder
+            ) + self.get_parameters(module.qf_twin)
+            optim_twin_critic = torch.optim.Adam(params_twin_critic, eps=1e-7)
+            self.register_optimizer(
+                module_id=module_id,
+                optimizer_name="qf_twin",
+                optimizer=optim_twin_critic,
+                params=params_twin_critic,
+                lr_or_lr_schedule=config.critic_lr,
+            )
+
+        # Define the optimizer for the actor.
+        params_actor = self.get_parameters(module.pi_encoder) + self.get_parameters(
+            module.pi
+        )
+        optim_actor = torch.optim.Adam(params_actor, eps=1e-7)
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="policy",
+            optimizer=optim_actor,
+            params=params_actor,
+            lr_or_lr_schedule=config.actor_lr,
+        )
+
+        # Define the optimizer for the value function.
+        params_value = self.get_parameters(module.vf_encoder) + self.get_parameters(
+            module.vf
+        )
+        optim_value = torch.optim.Adam(params_value, eps=1e-7)
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="value",
+            optimizer=optim_value,
+            params=params_value,
+            lr_or_lr_schedule=config.value_lr,
+        )
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: AlgorithmConfig,
+        batch: Dict,
+        fwd_out: Dict
+    ):
+
+        # Get the module and hyperparameters.
+        module = self._module[module_id]
+        expectile = self.expectile[module_id]
+        temperature = self.temperature[module_id]
+
+        # Get the action distribution for the actor loss.
+        action_train_dist_class = module.get_train_action_dist_cls()
+        action_train_dist = action_train_dist_class.from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+
+        # First, compute the value loss via the target Q-network and current observations.
+        value_loss = torch.mean(
+            self._expectile_loss(
+                fwd_out[QF_TARGET_PREDS] - fwd_out[Columns.VF_PREDS], expectile
+            )
+        )
+
+        # Second, compute the actor loss using the target-Q network and values.
+        exp_advantages = torch.minimum(
+            torch.exp(
+                temperature * (fwd_out[QF_TARGET_PREDS] - fwd_out[Columns.VF_PREDS])
+            ),
+            torch.Tensor([100.0]).to(self.device),
+        )
+        # Note, we are using here the actions from the data sample.
+        action_logps = action_train_dist.logp(batch[Columns.ACTIONS])
+        # Compute the actor loss.
+        actor_loss = -torch.mean(exp_advantages.detach() * action_logps)
+
+        # Third, compute the critic loss.
+        target_critic = (
+            batch[Columns.REWARDS]
+            + config.gamma
+            * (1 - batch[Columns.TERMINATEDS].float())
+            * fwd_out[VF_PREDS_NEXT].detach()
+        )
+
+        critic_loss = torch.mean(
+            torch.nn.MSELoss(reduction="none")(target_critic, fwd_out[QF_PREDS])
+        )
+
+        # If we have a twin-Q architecture, calculate the its loss, too.
+        if config.twin_q:
+            critic_twin_loss = (
+                torch.mean(
+                    torch.nn.MSELoss(reduction="none")(
+                        target_critic, fwd_out[QF_TWIN_PREDS]
+                    )
+                )
+                * 0.5
+            )
+            critic_loss *= 0.5
+
+        # Compute the total loss.
+        total_loss = value_loss + actor_loss + critic_loss
+
+        # If we have a twin-Q architecture, add its loss.
+        if config.twin_q:
+            total_loss += critic_twin_loss
+
+        # Log metrics.
+        self.metrics.log_dict(
+            {
+                POLICY_LOSS_KEY: actor_loss,
+                QF_LOSS_KEY: critic_loss,
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+
+        # Log the losses also in the temporary containers for gradient computation.
+        self._temp_losses[(module_id, POLICY_LOSS_KEY)] = actor_loss
+        self._temp_losses[(module_id, QF_LOSS_KEY)] = critic_loss
+        self._temp_losses[(module_id, VF_LOSS)] = value_loss
+
+        # If a twin-Q architecture is used add metrics and loss.
+        if config.twin_q:
+            self.metrics.log_value(
+                key=(module_id, QF_TWIN_LOSS_KEY),
+                value=critic_twin_loss,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+            self._temp_losses[(module_id, QF_TWIN_LOSS_KEY)] = critic_twin_loss
+
+        return total_loss
+
+    @override(TorchLearner)
+    def compute_gradients(
+        self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
+    ) -> ParamDict:
+        grads = {}
+        for module_id in set(loss_per_module.keys()) - {ALL_MODULES}:
+            # Loop through optimizers registered for this module.
+            for optim_name, optim in self.get_optimizers_for_module(module_id):
+                # Zero the gradients. Note, we need to reset the gradients b/c
+                # each component for a module operates on the same graph.
+                optim.zero_grad(set_to_none=True)
+
+                # Compute the gradients for the component and module.
+                loss_tensor = self._temp_losses.pop((module_id, optim_name + "_loss"))
+                loss_tensor.backward(retain_graph=True)
+                # Store the gradients for the component and module.
+                grads.update(
+                    {
+                        pid: p.grad
+                        for pid, p in self.filter_param_dict_for_optimizer(
+                            self._params, optim
+                        ).items()
+                    }
+                )
+
+        # Make sure we updated on all loss terms.
+        assert not self._temp_losses
+        return grads
+
+    def _expectile_loss(self, diff: TensorType, expectile: TensorType) -> TensorType:
+        """Computes the expectile loss.
+
+        Args:
+            diff: A tensor containing a difference loss.
+            expectile: The expectile to use for the expectile loss.
+
+        Returns:
+            The expectile loss of `diff` using `expectile`.
+        """
+        weight = torch.where(diff > 0, expectile, 1 - expectile)
+        return weight * torch.pow(diff, 2)