JIT accumulating GPU Memory

[felixmin]

Hey everybody!

In my training loop I have the following BCJaxPyPolicy class. The _action method is called during evaluation every x loops and then many times in a row.

However with every evaluation more memory is allocated, so that I eventually run out of memory. Does anybody know why this is the case?

class BCJaxPyPolicy(py_policy.PyPolicy):

  def __init__(self, ...):
    super(BCJaxPyPolicy, self).__init__(time_step_spec, action_spec)
    ...
    self._run_action_inference_jit = jax.jit(self._run_action_inference)

  def _run_action_inference(self, observation, goal, rng):
    # add batch dim
    observation = jax.tree_map(lambda x: jnp.expand_dims(x, 0), observation)
    goal = jax.tree_map(lambda x: jnp.expand_dims(x, 0), goal)

    normalized_action = self.sample_action(model=self.model, variables=self.variables,
                              observation=observation, goal=goal, rng=rng)

    action = (
        normalized_action * jnp.maximum(self.action_std, EPS) +
        self.action_mean)
    action = jnp.clip(action, self.action_spec.minimum, self.action_spec.maximum)
    return action

  def _action(self, time_step, policy_state=(), seed=0):
    ...
    action = self._run_action_inference_jit(observation, goal, rng)[0]
    ....