Mava icon indicating copy to clipboard operation
Mava copied to clipboard

[MAINTAIN] Separate functionality of tf and jax env loops

Open sash-a opened this issue 2 years ago • 0 comments

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

sash-a avatar Jun 27 '22 14:06 sash-a