Feat: sebulba rec ippo #1142
Feat: sebulba rec ippo #1142
Conversation
SimonDuToit
requested review from
RuanJohn,
sash-a,
OmaymaMahjoub,
WiemKhlifi and
Louay-Ben-nessir
as code owners
| @@ -0,0 +1,910 @@ | |||
| # Copyright 2022 InstaDeep Ltd. All rights reserved. | |||
There was a problem hiding this comment.
If you can update the typings in Pipeline mava/utils/sebulba.py to be Union[PPOTransition, RNNPPOTransition]
There was a problem hiding this comment.
This causes errors in the pre-commit. For now I changed both sebulba systems to use the MavaTransition type-var but this is probably a temporary solution.
There was a problem hiding this comment.
Can you please make an issue for this. I think the best solution is to make a protocol with the all the common things in a transition (actions, obs, done, reward). The challenge is that named tuples don't seem to work with protocols so we'd likely need to switch to a flax/chex dataclass
…nto feat/sebulba_rec_ippo
mava/systems/ppo/sebulba/rec_ippo.py
Outdated
| (config.arch.num_envs, num_agents), config.network.hidden_state_dim | ||
| ) | ||
| hstates = HiddenStates(init_policy_hstate, init_critic_hstate) | ||
| hstates_tpu = tree.map(move_to_device, hstates) |
There was a problem hiding this comment.
device_put already tree maps
| hstates_tpu = tree.map(move_to_device, hstates) | |
| hstates_tpu = move_to_device(hstates) |
mava/systems/ppo/sebulba/rec_ippo.py
Outdated
| obs_tpu = tree.map(move_to_device, timestep.observation) | ||
| last_dones = tree.map(move_to_device, dones) |
There was a problem hiding this comment.
| obs_tpu = tree.map(move_to_device, timestep.observation) | |
| last_dones = tree.map(move_to_device, dones) | |
| obs_tpu = move_to_device(timestep.observation) | |
| last_dones = move_to_device(dones) |
Sebulba implementation of recurrent IPPO.