Run PPO with task RotateSingleObjectInHandLevel0-v1

[Jianshu-Hu]

Hi,
Thanks again for maintaining such a good repo! I am trying to run PPO algorithm from https://github.com/haosulab/ManiSkill/blob/main/examples/baselines/ppo/ppo.py with the task RotateSingleObjectInHandLevel0-v1. However, I got this error in the first training iteration:

File "/bd_byt4090i1/users/jhu/SoftHandGrasping/train.py", line 166, in
next_obs, reward, terminations, truncations, infos = envs.step(clip_action(action))
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/mani_skill/vector/wrappers/gymnasium.py", line 152, in step
obs, infos = self.reset(options=dict(env_idx=env_idx))
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/mani_skill/vector/wrappers/gymnasium.py", line 93, in reset
obs, info = self._env.reset(seed=seed, options=options) # type: ignore
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/gymnasium/core.py", line 467, in reset
return self.env.reset(seed=seed, options=options)
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/gymnasium/wrappers/order_enforcing.py", line 61, in reset
return self.env.reset(**kwargs)
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/mani_skill/envs/sapien_env.py", line 947, in reset
self._initialize_episode(env_idx, options)
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py", line 159, in _initialize_episode
self._initialize_actors(env_idx)
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py", line 173, in _initialize_actors
pose.raw_pose[:, 0:3] = new_pos
File "/bd_byt4090i1/users/jhu/miniconda3/envs/ms3/lib/python3.10/site-packages/torch/utils/_device.py", line 106, in torch_function
return func(*args, **kwargs)
The expanded size of the tensor (512) must match the existing size (5) at non-singleton dimension 0. Target sizes: [512, 3]. Tensor sizes: [5, 3]

I found that in this function

ManiSkill/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py

Line 162 in 3bc2381

def _initialize_actors(self, env_idx: torch.Tensor):

, only partial environments (marked by env_idx, size 5) need to be reset while the function trying to set the object pose (size 512) of all environments. This is the cause of the error.

I am wondering:

1. In this initialization function, should all environments be reset or only environments marked with env_idx be reset?
2. If only environments marked with env_idx should be reset, after fixing this, I found the similar error happened in the evaluate function below

   ManiSkill/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py

   Line 229 in 3bc2381

   def evaluate(self, **kwargs) -> dict:

   ,
   and env_idx was not passed into this function. I have no idea how to fix this function to make it compatible with partial reset.

Thanks for any help in advance!

[StoneT2000]

This is one of the older maniskill envs that is a bit outdated before we started explicitly supporting partial resets. Happy to accept a PR to fix this, I unfortunately have very low bandwidth at the moment.

[Jianshu-Hu]

Ok, thanks for your reply! Could you indicate which environments are relatively new and support partial resets. I can try to fix this environment according to the examples.
And one thing I would like to double-check. For evaluate() function such as

ManiSkill/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py

Line 229 in 3bc2381

def evaluate(self, **kwargs) -> dict:

should it also be compatible with the partial reset so that only the terminated environments are evaluated?

[StoneT2000]

evaluate function is called at every timestep including during resets, so no it is not meant to do that. The functions with env_idx as arguments are usually the ones where you have to be careful, which is just the initialize episode function.