    def encode_images(
        self,
        images: jax.Array,
        rng: PRNGKey,
        *,
        grad_params: Optional[Params] = None,
        train: bool = True,
    ):
        """
        Forward pass for pre-trained encoder network.
        Pass grad_params to use non-default parameters (e.g. for gradients).
        """
        if train:
            assert rng is not None, "Must specify rng when training"
        return self.state.apply_fn(
            {"params": grad_params or self.state.params},
            images,
            name="vice",
            rngs={"dropout": rng} if train else {},
            return_encoded=True,
            train=train,
        )
Forward pass for pre-trained encoder network.
Pass grad_params to use non-default parameters (e.g. for gradients).
serl_launcher/serl_launcher/agents/continuous/vice.py