Help determining whether I'm using lax.scan and lax.cond correctly in re-writing existing JAX code to make it more efficient?

[Chulabhaya]

Hi! I have some JAX code for a reinforcement learning algorithm (SAC) that I initially worked on and is pretty fast compared to my existing PyTorch code, https://pastebin.com/tLqhSc6i. I figured that I could make this version of the code faster if I made everything end-to-end jittable and attempted to use lax.scan to replace the main for loop and lax.cond to replace the if statements as follows: https://pastebin.com/AbE1VTHh

My main question is did I implement this re-write correctly in terms of making the loops and conditionals also jit-compatible, and is it beneficial to code my algorithm in this way? I'm pretty new to JAX and still trying to learn my way around the best practices so any help would be greatly appreciated. Thank you in advance!

Everything is contained to one file in those scripts for the algorithms except for the replay buffer, which is the following:

from functools import partial

import chex
import jax


class UniformReplayBuffer:
    def __init__(
        self,
        buffer_size: int,
        batch_size: int,
    ) -> None:
        self.buffer_size = buffer_size
        self.batch_size = batch_size

    @partial(jax.jit, static_argnums=(0))
    def add(
        self,
        buffer_state: dict,
        experience: tuple,
        idx: int,
    ) -> dict:
        observation, action, reward, next_observation, done = experience
        idx = idx % self.buffer_size

        buffer_state["observations"] = (
            buffer_state["observations"].at[idx].set(observation)
        )
        buffer_state["actions"] = buffer_state["actions"].at[idx].set(action)
        buffer_state["rewards"] = buffer_state["rewards"].at[idx].set(reward)
        buffer_state["next_observations"] = (
            buffer_state["next_observations"].at[idx].set(next_observation)
        )
        buffer_state["dones"] = buffer_state["dones"].at[idx].set(done)

        return buffer_state

    @partial(jax.jit, static_argnums=(0))
    def sample(
        self,
        key: chex.PRNGKey,
        buffer_state: dict,
        current_buffer_size: int,
    ):

        @partial(jax.vmap, in_axes=(0, None))  # iterate over the indexes
        def sample_batch(indexes, buffer):
            """
            For a given index, extracts all the values from the buffer
            """
            return jax.tree_map(lambda x: x[indexes], buffer)

        key, sampling_key = jax.random.split(key)
        indexes = jax.random.randint(
            sampling_key,
            shape=(self.batch_size,),
            minval=0,
            maxval=current_buffer_size,
        )
        experiences = sample_batch(indexes, buffer_state)

        return experiences, key

[jakevdp]

Hi - I think you're going to have to narrow this down for us a bit. Your link is to a file with over 1000 lines of code, and it's not really clear which part you're asking about. I would suggest doing your best to create a minimal reproducible example of what you're attempting to do, and then use that to ask a more pointed question. Thanks!