A question about the sample_actions()

[fuyw]

implicit_q_learning/policy.py

Line 66 in 09d7002

@functools.partial(jax.jit, static_argnames=('actor_def', 'distribution'))

Hi Ilya,

Many thanks for the nice work. I have a question of the sample_actions() function, why do we need the _sample_actions()? Isn't it redundant?

Maybe we can simply:

@functools.partial(jax.jit, static_argnames=('actor_def'))
def sample_actions(rng, actor_def, actor_params, observations, temperature):
    dist = actor_def.apply({'params': actor_params}, observations, temperature)
    rng, key = jax.random.split(rng)
    return rng, dist.sample(seed=key)

Further, I tried to reimplement IQL with TrainState. I found that use TrainState is slower than this implementation (~100-200 fps).

[ikostrikov]

@fuyw I think there is some bug on Windows otherwise: ikostrikov/jaxrl#18

That's cool! Is it this implementation? I will take a look.

[fuyw]

Thanks for the reply. Yes it is, and I just refactored the code according to the flax official examples.

For simplicity, I replaced the tfd to distrax, and this does not matters in my experiments.

[fuyw]

Sorry Ilya, I found a bug in my previous implementation. I used a jax.device_put() when sampling from the buffer, which wastes time. When I fixed this bug, the throughput is close to this implementation now.