Significant slowdown when using jax.random.choice compared to np.random.choice

[dui1234]

Hi,

I have a function that return a batch of samples from dataset:

def sample_batch(self):
        idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
        return [self.obs_buf[idxs],
                self.next_obs_buf[idxs],
                self.acts_buf[idxs],
                self.rews_buf[idxs],
                self.done_buf[idxs],
                # for N-step Learning
                idxs
               ]
def sample_batch_from_idxs(self, idxs):
        return [self.obs_buf[idxs],
                self.next_obs_buf[idxs],
                self.acts_buf[idxs],
                self.rews_buf[idxs],
                self.done_buf[idxs]]

and this function is called under another function:

def get_samples(agent,s_params):
    samples1 = agent['memory'].sample_batch()
    indices = samples1[-1]
    samples2 = agent['memory_n'].sample_batch_from_idxs(indices)
    return [samples1,samples2,np.array(indices)]

This works perfectly find.
However, once I replace the idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
with jax.random.choice:

def sample_batch(self,key):
        #idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
        idxs = np.array(jax.random.choice(key,a=self.size,shape=(self.batch_size,),replace=False))
        return [self.obs_buf[idxs],
                self.next_obs_buf[idxs],
                self.acts_buf[idxs],
                self.rews_buf[idxs],
                self.done_buf[idxs],
                # for N-step Learning
                idxs
               ]

and slightly change the get_samples function as:

def get_samples(agent,s_params,key):
    key,sp =  jax.random.split(key,2)
    samples1 = agent['memory'].sample_batch(sp)
    indices = samples1[-1]
    samples2 = agent['memory_n'].sample_batch_from_idxs(indices)
    return [samples1,samples2,np.array(indices)],key

The speed of the entile code is sharply dropped. Just wondering what should I do in this case?

[soraros]

You should jit your call t jax.random.choice.

from functools import partial

@partial(jit, static_argnums=(1, 2))
def choice(key, a: int, shape: tuple[int, ...]):
  return jax.random.choice(key, a=a, shape=shape, replace=False)

def sample_batch(self, key):
  # pretty sure you need conversion to np.ndarray here
  idxs = np.array(choice(key, a=self.size, shape=(self.batch_size,)))
  return [
    self.obs_buf[idxs],
    self.next_obs_buf[idxs],
    self.acts_buf[idxs],
    self.rews_buf[idxs],
    self.done_buf[idxs],
    # for N-step Learning
    idxs
  ]

BTW, it would help people help you better if you could provide a minimal reproducible example.