    def update_vice(
        self,
        batch,
        pmap_axis: Optional[str] = None,
    ):
        """
        update the VICE reward classifier using the BCE loss.
        addtional regularization techniques are also used: mixup, label smoothing, and gradient penalty regularization
        to prevent GAN mode collapse.

        NOTE: assumes that the second half of the batch contains the goal images, so labels = 1
        """
        new_agent = self
        rng = new_agent.state.rng
        if self.config["image_keys"][0] not in batch["next_observations"]:
            batch = _unpack(batch)

        def mixup_data_rng(key_0, key_1, x: jnp.ndarray, y: jnp.ndarray, alpha=1):
            """
            performs mixup regularization on the input images and labels
            """
            if alpha > 0:
                lam = jax.random.beta(key_0, alpha, alpha)
            else:
                lam = 1
            batch_size = x.shape[0]
            index = jax.random.permutation(key_1, batch_size)
            mixed_x = lam * x + (1 - lam) * x[index, :]
            y_a, y_b = y, y[index]
            return mixed_x, y_a, y_b, lam

        observations = batch["next_observations"]
        key, rng = jax.random.split(rng)
        # data augmentation on images
        aug_observations = self.data_augmentation_fn(key, observations)
        mix_observations = observations.unfreeze()
        gp_observations = observations.unfreeze()
        all_info = {}

        for image_key in self.config["image_keys"]:
            pixels = observations[image_key]
            batch_size = pixels.shape[0]

            pixels = observations[image_key][: batch_size // 2]
            aug_pixels = aug_observations[image_key][: batch_size // 2]
            goal_pixels = observations[image_key][batch_size // 2 :]
            aug_goal_pixels = aug_observations[image_key][batch_size // 2 :]

            # concatenate all images for update
            all_obs_pixels = jnp.concatenate([pixels, aug_pixels], axis=0)
            all_goal_pixels = jnp.concatenate([goal_pixels, aug_goal_pixels], axis=0)
            all_pixels = jnp.concatenate([all_goal_pixels, all_obs_pixels], axis=0)

            # create labels
            ones = jnp.ones((batch_size, 1))
            zeros = jnp.zeros((batch_size, 1))
            y_batch = jnp.concatenate([ones, zeros], axis=0)
            y_batch = y_batch.squeeze(-1)

            # label smoothing, help with bad numerical issue with large negative logits
            y_batch = y_batch * (1 - 0.2) + 0.5 * 0.2

            # encode images into embeddings
            key, rng = jax.random.split(rng)
            encoded = self.encode_images(all_pixels, key, train=True)

            # perform mixup
            key_0, key_1, rng = jax.random.split(rng, 3)
            mix_encoded, y_a_0, y_b_0, lam_0 = mixup_data_rng(
                key_0, key_1, encoded, y_batch
            )
            mix_observations[image_key] = mix_encoded

            # interpolate for Gradient Penalty regularization
            key, rng = jax.random.split(rng)
            """
            generate random epsilon for each sample in the batch, the shape here depends on the shape of the encoded embeddings.
            Here are some examples:
            epsilon = jax.random.uniform(key, shape=(all_obs_pixels.shape[0], 1, 1, 1, 1))
            epsilon = jax.random.uniform(key, shape=(all_obs_pixels.shape[0], 1))
            """
            # epsilon = jax.random.uniform(key, shape=(all_obs_pixels.shape[0], 1, 1, 1))
            epsilon = jax.random.uniform(
                key,
                shape=(len(mix_encoded) // 2, *([1] * (len(mix_encoded.shape[1:])))),
            )
            gp_encoded = (
                epsilon * mix_encoded[: len(mix_encoded) // 2]
                + (1 - epsilon) * mix_encoded[len(mix_encoded) // 2 :]
            )
            gp_observations[image_key] = gp_encoded

        # remove all non pixel inputs keys from the batch
        remove_keys = [
            k for k in gp_observations.keys() if k not in self.config["image_keys"]
        ]
        for k in remove_keys:
            gp_observations.pop(k)

        mix_observations = frozen_dict.freeze(mix_observations)
        gp_observations = frozen_dict.freeze(gp_observations)

        key, rng = jax.random.split(rng, 2)

        def mixup_loss_fn(params, rng) -> Tuple[jnp.ndarray, Dict[str, float]]:
            y_hat = new_agent.state.apply_fn(
                {
                    "params": params,
                },
                mix_observations,
                name="vice",
                train=True,
                classify_encoded=True,
                rngs={"dropout": rng},
            )
            bce_loss_a = jnp.mean(optax.sigmoid_binary_cross_entropy(y_hat, y_a_0))
            bce_loss_b = jnp.mean(optax.sigmoid_binary_cross_entropy(y_hat, y_b_0))
            bce_loss = lam_0 * bce_loss_a + (1 - lam_0) * bce_loss_b
            return bce_loss, {
                "bce_loss": bce_loss,
            }

        def gp_loss_fn(params, rng) -> Tuple[jnp.ndarray, Dict[str, float]]:
            bce_loss, info = mixup_loss_fn(params, key)
            helper_fn = lambda x: new_agent.state.apply_fn(
                {
                    "params": params,
                },
                x,
                name="vice",
                train=True,
                classify_encoded=True,
                rngs={"dropout": rng},
            )

            grad_wrt_input = jax.vmap(jax.grad(helper_fn), in_axes=0, out_axes=0)
            gradients = grad_wrt_input(gp_observations)
            gradients = jnp.concatenate([grads for grads in gradients.values()], axis=0)
            gradients = gradients.reshape((gradients.shape[0], -1))
            grad_norms = jnp.sqrt(jnp.sum((gradients**2 + 1e-6), axis=1))
            grad_penalty = jnp.mean((grad_norms - 1) ** 2)

            return bce_loss + 10 * grad_penalty, {
                "bce_loss": bce_loss,
                "grad_norm": grad_norms.mean(),
            }

        loss_fns = {
            "actor": lambda params, rng: (0.0, {}),
            "critic": lambda params, rng: (0.0, {}),
            "temperature": lambda params, rng: (0.0, {}),
            "vice": gp_loss_fn,
        }
        new_state, info = new_agent.state.apply_loss_fns(
            loss_fns, pmap_axis=pmap_axis, has_aux=True
        )
        new_state = new_state.replace(rng=rng)
        all_info.update(info)

        return new_agent.replace(state=new_state), all_info
update the VICE reward classifier using the BCE loss.
addtional regularization techniques are also used: mixup, label smoothing, and gradient penalty regularization
to prevent GAN mode collapse.

NOTE: assumes that the second half of the batch contains the goal images, so labels = 1
serl_launcher/serl_launcher/agents/continuous/vice.py