    def update_high_utd(
        self,
        batch: Batch,
        *,
        utd_ratio: int,
        pmap_axis: Optional[str] = None,
    ) -> Tuple["DrQAgent", dict]:
        """
        Fast JITted high-UTD version of `.update`.

        Splits the batch into minibatches, performs `utd_ratio` critic
        (and target) updates, and then one actor/temperature update.

        Batch dimension must be divisible by `utd_ratio`.

        It also performs data augmentation on the observations and next_observations
        before updating the network.
        """
        new_agent = self
        if self.config["image_keys"][0] not in batch["next_observations"]:
            batch = _unpack(batch)

        rng = new_agent.state.rng
        rng, obs_rng, next_obs_rng = jax.random.split(rng, 3)
        obs = self.data_augmentation_fn(obs_rng, batch["observations"])
        next_obs = self.data_augmentation_fn(next_obs_rng, batch["next_observations"])
        rewards = (self.vice_reward(next_obs) >= 0.5) * 1.0
        batch = batch.copy(
            add_or_replace={
                "observations": obs,
                "next_observations": next_obs,
                "rewards": rewards,
            }
        )

        new_state = self.state.replace(rng=rng)

        new_agent = self.replace(state=new_state)
        new_agnet, info = SACAgent.update_high_utd(
            new_agent, batch, utd_ratio=utd_ratio, pmap_axis=pmap_axis
        )
        info["vice_rewards"] = rewards.mean()
        return new_agent, info
Fast JITted high-UTD version of `.update`.

Splits the batch into minibatches, performs `utd_ratio` critic
(and target) updates, and then one actor/temperature update.

Batch dimension must be divisible by `utd_ratio`.

It also performs data augmentation on the observations and next_observations
before updating the network.
serl_launcher/serl_launcher/agents/continuous/vice.py