Feat Sebulba recurrent IQL #1148
Conversation
Louay-Ben-nessir
requested review from
RuanJohn,
sash-a,
OmaymaMahjoub,
WiemKhlifi and
SimonDuToit
as code owners
|
|
||
| # todo: remove the ppo dependencies when we make sebulba for other systems |
There was a problem hiding this comment.
This is a good point though, maybe there's something we can do about it 🤔
Maybe a protocol like that has action, obs, reward, not sure if there's any other common attributes?
| terminated = np.repeat( | ||
| terminated[..., np.newaxis], repeats=self.num_agents, axis=-1 | ||
| ) # (B,) --> (B, N) |
There was a problem hiding this comment.
Does this already happen for smax and lbf?
| target: Array, | ||
| ) -> Tuple[Array, Metrics]: | ||
| # axes switched here to scan over time | ||
| hidden_state, obs_term_or_trunc = prep_inputs_to_scannedrnn(obs, term_or_trunc) |
There was a problem hiding this comment.
A general comment, I think this would be a lot easier to read if we used done to mean term_or_trunc which I think is a reasonable thing. Would have to make the change in anakin also though :/
| """ | ||
|
|
||
| eps = jnp.maximum( | ||
| config.system.eps_min, 1 - (t / config.system.eps_decay) * (1 - config.system.eps_min) |
There was a problem hiding this comment.
Would be nice if we could set a different decay per actor, although I think that's out of scope for this PR. Maybe if you could make an issue to add in some of the ape-X DQN features that would be great
There was a problem hiding this comment.
I can easily add this in this PR 👀
There was a problem hiding this comment.
I think rather leave it for now, no need to make this more complex
| rewards = np.zeros((num_envs, num_agents), dtype=float) | ||
| teminated = np.zeros(num_envs, dtype=float) | ||
| rewards = np.zeros((num_envs, self.num_agents), dtype=float) | ||
| terminated = np.zeros((num_envs, self.num_agents), dtype=float) |
There was a problem hiding this comment.
I assume with this change we would also need to change sebulba PPO? Since currently it does this same operation. We should decide if generally its better doing this in the system or the wrapper.
There was a problem hiding this comment.
Good point, I think we do this in the wrappers for the anakin systems
What?
A recurrent IQL implementation using the Sebulba architecture.
Why?
Offline Sebulba base and non-jax envs in Mava.
How?
Mixed the Sebulba structure from PPO with the learner code from Anakin IQL.