from typing import Any, Union
import numpy as np
import sapien
import torch
from mani_skill.agents.robots import Fetch, Panda
from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.sensors.camera import CameraConfig
from mani_skill.utils import common, sapien_utils
from mani_skill.utils.registration import register_env
from mani_skill.utils.structs.types import SimConfig
@register_env("CustomEnv-v1", max_episode_steps=200)
[docs]class CustomEnv(BaseEnv):
[docs] SUPPORTED_ROBOTS = ["panda", "fetch"]
[docs] agent: Union[Panda, Fetch]
def __init__(self, *args, robot_uids="panda", robot_init_qpos_noise=0.02, **kwargs):
[docs] self.robot_init_qpos_noise = robot_init_qpos_noise
super().__init__(*args, robot_uids=robot_uids, **kwargs)
@property
[docs] def _default_sim_config(self):
return SimConfig()
@property
[docs] def _default_sensor_configs(self):
pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
return [
CameraConfig("base_camera", pose=pose, width=128, height=128, fov=np.pi / 2)
]
@property
[docs] def _default_human_render_camera_configs(self):
pose = sapien_utils.look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
return CameraConfig("render_camera", pose=pose, width=512, height=512, fov=1)
[docs] def _load_agent(self, options: dict):
super()._load_agent(options, sapien.Pose(p=[0, 0, 0]))
[docs] def _load_scene(self, options: dict):
pass
[docs] def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
pass
[docs] def evaluate(self):
return {
"success": torch.zeros(self.num_envs, device=self.device, dtype=bool),
"fail": torch.zeros(self.num_envs, device=self.device, dtype=bool),
}
[docs] def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
return torch.zeros(self.num_envs, device=self.device)
[docs] def compute_normalized_dense_reward(
self, obs: Any, action: torch.Tensor, info: dict
):
max_reward = 1.0
return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
