dev: brax automatic detect batch size

Author: BillHuang2001

src/evox/problems/neuroevolution/reinforcement_learning/brax.py

@@ -4,6 +4,7 @@
 import jax
 from jax import jit, vmap
 import jax.numpy as jnp
+from jax.tree_util import tree_leaves
 from evox import Problem, State, jit_method
 
 
@@ -12,7 +13,6 @@ def __init__(
         self,
         policy: Callable,
         env_name: str,
-        batch_size: int,
         cap_episode: int,
         backend: str = "generalized",
     ):
@@ -41,17 +41,17 @@ def __init__(
         self.env = envs.wrappers.training.VmapWrapper(
             envs.get_environment(env_name=env_name, backend=backend)
         )
-        self.batch_size = batch_size
         self.cap_episode = cap_episode
         self.jit_reset = jit(self.env.reset)
         self.jit_env_step = jit(self.env.step)
 
     def setup(self, key):
-        return State(init_state=self.jit_reset(jnp.tile(key, (self.batch_size, 1))))
+        return State(key=key)
 
     @jit_method
     def evaluate(self, state, weights):
-        brax_state = state.init_state
+        batch_size = tree_leaves(weights)[0].shape[0]
+        brax_state = self.jit_reset(jnp.tile(state.key, (batch_size, 1)))
 
         def cond_func(val):
             counter, state, _total_reward = val
@@ -64,7 +64,7 @@ def body_func(val):
             total_reward += (1 - brax_state.done) * brax_state.reward
             return counter + 1, brax_state, total_reward
 
-        init_val = (0, brax_state, jnp.zeros((self.batch_size,)))
+        init_val = (0, brax_state, jnp.zeros((batch_size,)))
 
         _counter, _brax_state, total_reward = jax.lax.while_loop(
             cond_func, body_func, init_val
@@ -73,7 +73,14 @@ def body_func(val):
         return total_reward, state
 
     def visualize(
-        self, state, key, weights, output_type: str = "HTML", *args, **kwargs
+        self,
+        state,
+        key,
+        weights,
+        output_type: str = "HTML",
+        respect_done=False,
+        *args,
+        **kwargs,
     ):
         env = envs.get_environment(env_name=self.env_name, backend=self.backend)
         brax_state = jax.jit(env.reset)(key)
@@ -86,7 +93,7 @@ def visualize(
             trajectory.append(brax_state.pipeline_state)
             episode_length += 1 - brax_state.done
 
-            if brax_state.done:
+            if respect_done and brax_state.done:
                 break
 
         if output_type == "HTML":