Repository: EMI-Group/evox

Issue:
Encounter Errors when Using HPO with EvoX  involves Brax problems
   Hello!  When I attempt to use HPO with EvoX, specifically when the inner problem involves Brax problems, the implementation does not work as expected. I have tried to debug the issue but have not been able to resolve it.
The current examples in the documentation only demonstrate how to deploy HPO when the inner problems are numerical optimization problems. Could you please provide an example or guidance on how to implement HPO with EvoX for Brax problems?
Additionally, I encountered the following error during my attempts:

```
 File "/mnt/d/pycharm_project/pythonProject/evox/src/evox/workflows/std_workflow.py", line 109, in _evaluate
    fitness = self.problem.evaluate(population)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/d/pycharm_project/pythonProject/evox/src/evox/problems/neuroevolution/brax.py", line 203, in evaluate
    model_state, rewards = self._evaluate_brax(model_state)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 745, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/mnt/d/pycharm_project/pythonProject/evox/src/evox/problems/neuroevolution/brax.py", line 177, in _evaluate_brax
    brax_state = self.vmap_brax_step(brax_state, to_jax_array(action))
                                                 ^^^^^^^^^^^^^^^^^^^^
  File "/mnt/d/pycharm_project/pythonProject/evox/src/evox/problems/neuroevolution/brax.py", line 25, in to_jax_array
    return jax.dlpack.from_dlpack(x.detach())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/jax/_src/dlpack.py", line 276, in from_dlpack
    return _from_dlpack(external_array, device, copy)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/jax/_src/dlpack.py", line 228, in _from_dlpack
    dlpack = external_array.__dlpack__(stream=stream)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/_tensor.py", line 1675, in __dlpack__
    return handle_torch_function(Tensor.__dlpack__, (self,), self, stream)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/overrides.py", line 1720, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/d/pycharm_project/pythonProject/evox/src/evox/core/module.py", line 353, in __torch_function__
    return func(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/_tensor.py", line 1675, in __dlpack__
    return handle_torch_function(Tensor.__dlpack__, (self,), self, stream)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/overrides.py", line 1720, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/utils/_device.py", line 104, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/ding/anaconda3/envs/torch_env/lib/python3.11/site-packages/torch/_tensor.py", line 1724, in __dlpack__
    return torch.to_dlpack(self)
           ^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
ERROR conda.cli.main_run:execute(125): `conda run python /mnt/d/pycharm_project/pythonProject/evox/src/te.py` failed. (See above for error)
```

Hints from issue comments:
@BillHuang2001: Thank you for reporting this issue. The root cause is that EvoX uses PyTorch, while Brax relies on JAX, requiring tensor conversions between the two frameworks. And when using HPO, we utilize torch.vmap, but currently, torch.vmap does not support DLPack conversion.

We have forwarded this issue to the upstream. (https://github.com/pytorch/pytorch/issues/148306)@BillHuang2001: By the way, if you're looking to build a Population-Based Training (PBT) workflow or RL-related workflow, consider trying out [EvoRL](https://github.com/EMI-Group/evorl), our sister project.

Solution patch:
diff --git a/src/evox/problems/neuroevolution/brax.py b/src/evox/problems/neuroevolution/brax.py
index 86d7e27b..469cbe36 100644
--- a/src/evox/problems/neuroevolution/brax.py
+++ b/src/evox/problems/neuroevolution/brax.py
@@ -11,6 +11,7 @@ import torch.nn as nn
 import torch.utils.dlpack
 from brax import envs
 from brax.io import html, image
+from torch._C._functorch import get_unwrapped, is_batchedtensor
 
 from evox.core import Problem, use_state
 from evox.utils import VmapInfo
@@ -22,6 +23,8 @@ from .utils import get_vmap_model_state_forward
 # because they have a __dlpack__ method, which is called by their respective from_dlpack methods.
 def to_jax_array(x: torch.Tensor) -> jax.Array:
     # When the torch has GPU support but the jax does not, we need to move the tensor to CPU first.
+    if is_batchedtensor(x):
+        x = get_unwrapped(x)
     if x.device.type != "cpu" and jax.default_backend() == "cpu":
         return jax.dlpack.from_dlpack(x.detach().cpu())
     return jax.dlpack.from_dlpack(x.detach())
diff --git a/src/evox/workflows/eval_monitor.py b/src/evox/workflows/eval_monitor.py
index 4bc3ea66..806a987b 100644
--- a/src/evox/workflows/eval_monitor.py
+++ b/src/evox/workflows/eval_monitor.py
@@ -2,6 +2,7 @@ import warnings
 from typing import Dict, List, Tuple
 
 import torch
+from torch._C._functorch import get_unwrapped, is_batchedtensor
 
 from evox.core import Monitor, Mutable
 from evox.operators.selection import non_dominate_rank
@@ -105,8 +106,8 @@ class EvalMonitor(Monitor):
                 topk_solutions = torch.concatenate([self.topk_solutions, self.latest_solution])
                 topk_fitness = torch.concatenate([self.topk_fitness, fitness])
                 rank = torch.topk(topk_fitness, self.topk, largest=False)[1]
-                self.topk_fitness.copy_(topk_fitness[rank])
-                self.topk_solutions.copy_(topk_solutions[rank])
+                self.topk_fitness = topk_fitness[rank]
+                self.topk_solutions = topk_solutions[rank]
         elif fitness.ndim == 2:
             # multi-objective
             self.multi_obj = True
@@ -122,9 +123,15 @@ class EvalMonitor(Monitor):
     def record_history(self):
         if self.full_fit_history or self.full_sol_history:
             if self.full_sol_history:
-                self.solution_history.append(self.latest_solution.to(self.device))
+                latest_solution = self.latest_solution.to(self.device)
+                if is_batchedtensor(self.latest_solution):
+                    latest_solution = get_unwrapped(latest_solution)
+                self.solution_history.append(latest_solution)
             if self.full_fit_history:
-                self.fitness_history.append(self.latest_fitness.to(self.device))
+                latest_fitness = self.latest_fitness.to(self.device)
+                if is_batchedtensor(self.latest_fitness):
+                    latest_fitness = get_unwrapped(latest_fitness)
+                self.fitness_history.append(latest_fitness)
 
     def get_latest_fitness(self) -> torch.Tensor:
         """Get the fitness values from the latest iteration."""


Test patch:
diff --git a/unit_test/workflows/test_std_workflow.py b/unit_test/workflows/test_std_workflow.py
index 9db66acc..731a5b7b 100644
--- a/unit_test/workflows/test_std_workflow.py
+++ b/unit_test/workflows/test_std_workflow.py
@@ -34,8 +34,8 @@ class BasicAlgorithm(Algorithm):
     def step(self):
         pop = torch.rand(self.pop_size, self.dim, dtype=self.lb.dtype, device=self.lb.device)
         pop = pop * (self.ub - self.lb)[None, :] + self.lb[None, :]
-        self.pop.copy_(pop)
-        self.fit.copy_(self.evaluate(pop))
+        self.pop = pop
+        self.fit = self.evaluate(pop)
 
 
 class TestStdWorkflow(unittest.TestCase):
@@ -43,7 +43,8 @@ class TestStdWorkflow(unittest.TestCase):
         torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
         self.algo = BasicAlgorithm(10, -10 * torch.ones(2), 10 * torch.ones(2))
         self.prob = BasicProblem()
-        self.workflow = StdWorkflow(self.algo, self.prob)
+        self.monitor = EvalMonitor()
+        self.workflow = StdWorkflow(self.algo, self.prob, monitor=self.monitor)
 
     def test_basic_workflow(self):
         compiled_step = torch.compile(self.workflow.step)
@@ -58,6 +59,8 @@ class TestStdWorkflow(unittest.TestCase):
         self.assertIsNotNone(state)
         compiled_state_step = torch.compile(vmap_state_step)
         self.assertIsNotNone(compiled_state_step(state))
+        # Test if the monitor's history is a normal tensor (not batched)
+        self.assertIsNotNone(self.monitor.fitness_history[0] * 1)
 
     def test_classic_workflow(self):
         class solution_transform(nn.Module):