(self, step, worker=0)

(self, step)

(self, index, length, mapping)

(self)

(self, init_policy=None)

(self, key, uprob)

(self, callback)

(self, x)

(self, generator)

(self, action)

(self)

(dataset)

(self, key)

(self, action)

(self, env, key, value, space)

(param, update)

(name)

(self)

(name, df, fn, baseline=None)

(self, data)

(self, logits)

(self, child)

(config, replay, mode)

(self)

(self, st)

(self, counter, sharding)

(self, data=None, directory=None, amount=None)

(self, seed, shape=())

(self)

(self)

(self)

(x, fwd=None, bwd=None)

(self)

(self, carry, obs, prevact, training)

(self, var, x)

(xs, mask)

(sequence)

(self, carry, obs, mode)

(fn)

(self, tree)

(self, outs)

(self, batch_size)

(self, data)

(config)

(config)

(self)

(self)

(self, action)

(self, data)

(self, *args, **kwargs)

(self, nest, prefix=None)

(self, *args, method=None, **kwargs)

(self, spaces, batch_shape)

(name)

( ax, xs, ys, lo=None, hi=None, label=None, order=None, color=None, scatter=True, **kwargs)

(data)

(self, batch_size)

(self, bsize)

()

(self, other)

(self, carry, tokens, action, reset, training, single=False)

(source, prefetch=2)

(self)

(self, carry, data)

(self, index, length)

(self, carry, data)

(self, entries, carry=None)

(self, model, obs_space, act_space, config, jaxcfg)

(self, key)

(self, carry, data)

(self, value, space, key)

(self, chunk, worker)

(self, spaces)

(self, space)

(self, space)

(self, deter, stoch, action)

(self, tree)

(self, compiled)

(self, chunkid, index, keys=None, concat=True)

(self)

(self, chunkid, index)

(self, name, x)

(self, length=10)

(self, x)

(self, chunks)

(self, reward, is_first=False, is_last=False, is_terminal=False)

( self, image, reward, info, is_first=False, is_last=False, is_terminal=False)

( self, obs, reward, is_first=False, is_last=False, is_terminal=False)

(self, reward, is_first=False, is_last=False, is_terminal=False)

(self, reward, is_first=False, is_last=False, is_terminal=False)

(self, reward, is_first=False, is_last=False)

(self, carry, tokens, action, reset, training)

(self, index)

(self, x, q)

(self, feat)

(self, pipe)

(self, reset=False)

(self)

(self, image, size, method)

(self, shape, dtype)