Mava
Mava copied to clipboard
[MAINTAIN] Separate functionality of tf and jax env loops
Please describe what needs to be maintained?
This is really messy and is all over the show in environment loops and loop wrappers, so it would be really nice to have two separate env loops for tf and jax mava in order to get rid of this.
# Record counts.
if hasattr(self._executor, "_counts"):
loop_type = "evaluator" if "_" not in self._loop_label else "executor"
if hasattr(self._executor, "_variable_client"):
self._executor._variable_client.add_async(
[f"{loop_type}_episodes", f"{loop_type}_steps"],
{f"{loop_type}_episodes": 1, f"{loop_type}_steps": episode_steps},
)
else:
self._executor._counts[f"{loop_type}_episodes"] += 1
self._executor._counts[f"{loop_type}_steps"] += episode_steps
counts = self._executor._counts
elif hasattr(self._executor, "store"):
loop_type = "evaluator" if "_" not in self._loop_label else "executor"
self._executor.store.executor_parameter_client.add_async(
{f"{loop_type}_episodes": 1, f"{loop_type}_steps": episode_steps}
)
counts = self._executor.store.executor_counts
else:
counts = self._counter.increment(episodes=1, steps=episode_steps)
Describe the outcome you'd like
No hasattrs
in the loop wrapper
How do we know when maintenance is complete?
Checklist:
- [ ] there is a separate tf and mava environment loop
- [ ] there are no
hasattrs
in the code