fix(wzl): fix bugs for dict obs in new pipeline

ding/envs/env_manager/subprocess_env_manager.py

@@ -940,8 +940,7 @@ def ready_obs(self) -> tnp.array:
                 )
             time.sleep(0.001)
             sleep_count += 1
-        obs = [self._ready_obs[i] for i in self.ready_env]
-        return tnp.stack(obs)
+        return tnp.stack([tnp.array(self._ready_obs[i]) for i in self.ready_env])
 
     def step(self, actions: List[tnp.ndarray]) -> List[tnp.ndarray]:
         """

ding/framework/middleware/functional/collector.py

@@ -4,7 +4,7 @@
 import treetensor.torch as ttorch
 from ding.envs import BaseEnvManager
 from ding.policy import Policy
-from ding.torch_utils import to_ndarray
+from ding.torch_utils import to_ndarray, get_shape0
 
 if TYPE_CHECKING:
     from ding.framework import OnlineRLContext
@@ -74,8 +74,8 @@ def _inference(ctx: "OnlineRLContext"):
         obs = ttorch.as_tensor(env.ready_obs).to(dtype=ttorch.float32)
         ctx.obs = obs
         # TODO mask necessary rollout
-
-        obs = {i: obs[i] for i in range(obs.shape[0])}  # TBD
+        num_envs = get_shape0(obs)
+        obs = {i: obs[i] for i in range(num_envs)}  # TBD
         inference_output = policy.forward(obs, **ctx.collect_kwargs)
         ctx.action = [to_ndarray(v['action']) for v in inference_output.values()]  # TBD
         ctx.inference_output = inference_output

ding/framework/middleware/functional/evaluator.py

@@ -12,7 +12,7 @@
 from ding.policy import Policy
 from ding.data import Dataset, DataLoader
 from ding.framework import task
-from ding.torch_utils import to_list, to_ndarray
+from ding.torch_utils import to_list, to_ndarray, get_shape0
 from ding.utils import lists_to_dicts
 
 if TYPE_CHECKING:
@@ -251,7 +251,8 @@ def _evaluate(ctx: "OnlineRLContext"):
 
         while not eval_monitor.is_finished():
             obs = ttorch.as_tensor(env.ready_obs).to(dtype=ttorch.float32)
-            obs = {i: obs[i] for i in range(obs.shape[0])}  # TBD
+            num_envs = get_shape0(obs)
+            obs = {i: obs[i] for i in range(num_envs)}  # TBD
             inference_output = policy.forward(obs)
             if render:
                 eval_monitor.update_video(env.ready_imgs)

ding/torch_utils/__init__.py

@@ -1,6 +1,6 @@
 from .checkpoint_helper import build_checkpoint_helper, CountVar, auto_checkpoint
 from .data_helper import to_device, to_tensor, to_ndarray, to_list, to_dtype, same_shape, tensor_to_list, \
-    build_log_buffer, CudaFetcher, get_tensor_data, unsqueeze, squeeze, get_null_data
+    build_log_buffer, CudaFetcher, get_tensor_data, unsqueeze, squeeze, get_null_data, get_shape0
 from .distribution import CategoricalPd, CategoricalPdPytorch
 from .metric import levenshtein_distance, hamming_distance
 from .network import *

ding/torch_utils/data_helper.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import torch
+import treetensor.torch as ttorch
 
 
 def to_device(item: Any, device: str, ignore_keys: list = []) -> Any:
@@ -408,3 +409,20 @@ def get_null_data(template: Any, num: int) -> List[Any]:
         data['reward'].zero_()
         ret.append(data)
     return ret
+
+
+def get_shape0(data):
+    if isinstance(data, torch.Tensor):
+        return data.shape[0]
+    elif isinstance(data, ttorch.Tensor):
+
+        def fn(t):
+            item = list(t.values())[0]
+            if np.isscalar(item[0]):
+                return item[0]
+            else:
+                return fn(item)
+
+        return fn(data.shape)
+    else:
+        raise TypeError("not support type: {}".format(data))