Question about policy(...) implementation.

[edwhu]

I am interested in adding more behaviors.

In the policy function of agent.py, all behaviors are called before the mode is even checked.

dreamerv3/dreamerv3/agent.py

Lines 51 to 60 in 423291a

	def policy(self, obs, state, mode='train'):
	self.config.jax.jit and print('Tracing policy function.')
	obs = self.preprocess(obs)
	(prev_latent, prev_action), task_state, expl_state = state
	embed = self.wm.encoder(obs)
	latent, _ = self.wm.rssm.obs_step(
	prev_latent, prev_action, embed, obs['is_first'])
	self.expl_behavior.policy(latent, expl_state)
	task_outs, task_state = self.task_behavior.policy(latent, task_state)
	expl_outs, expl_state = self.expl_behavior.policy(latent, expl_state)

This doesn't seem like it will scale if we add a lot of different behaviors.

Is this because we want Jax to trace all variables in the train function, i.e. in the JaxAgent where we get the initial variables of the train function for optimization?

dreamerv3/dreamerv3/jaxagent.py

Line 228 in 423291a

varibs = self._train(varibs, rng, data, state, init_only=True)

[danijar]

Yes, that's the reason

[edwhu]

Thanks. It seems straightforward to replace the if statements with jax.lax.select or jax.lax.cond logic and move the policy calls into the branches - any particular reason why you chose to use python if statements instead?

[danijar]

A cond has runtime overhead (and also wouldn't work with the current API because JAX/GPUs don't support string types: jax-ml/jax#3045).