PickCubeEnv

class mani_skill.envs.tasks.tabletop.PickCubeEnv(robot_uids='panda', robot_init_qpos_noise=0.02, **kwargs)

A simple manipulation task where the objective is to grasp a red cube with a robot arm and move it to a target goal position. This environment serves as a baseline task to test whether a robot with manipulation capabilities can be simulated and trained properly.

Task Description

The agent must:

1. Grasp a red cube placed on a table
2. Move the cube to a target goal position (marked by a green sphere)
3. Place the cube within the goal threshold distance
4. Keep the robot static after placement

Randomizations

Each episode, the following are randomized:

• Cube position: The cube’s xy position is randomized on top of the table in the region [0.1, 0.1] x [-0.1, -0.1]. It is placed flat on the table.
• Cube rotation: The cube’s z-axis rotation is randomized to a random angle.
• Goal position: The target goal position (marked by a green sphere) has its xy position randomized in the region [0.1, 0.1] x [-0.1, -0.1] and z randomized in [0, 0.3].

Success Conditions

The task is considered successful when:

• The cube position is within goal_thresh (default 0.025m) euclidean distance of the goal position
• The robot is static (joint velocity < 0.2)

Args

• robot_uids (str, optional): Robot to use. Must be one of ["panda", "fetch", "xarm6_robotiq", "so100", "widowxai"]. Defaults to "panda".
• robot_init_qpos_noise (float, optional): Noise added to the initial robot joint positions. Defaults to 0.02.
• **kwargs: Additional arguments passed to BaseEnv. See BaseEnv documentation for details.

Attributes

• agent (Union[Panda, Fetch, XArm6Robotiq, SO100, WidowXAI]): The robot agent instance.
• cube: Actor - 可移动方块. [blog]
• goal_thresh (float): Distance threshold (in meters) for considering the cube successfully placed. Defaults to 0.025.
• cube_half_size (float): Half-size of the cube in meters. Defaults to 0.02.
• cube_spawn_half_size (float): Half-size of the cube spawn region. Defaults to 0.05.
• cube_spawn_center (tuple): Center of the cube spawn region. Defaults to (0, 0).
• SUPPORTED_ROBOTS (List[str]): List of supported robot UIDs: ["panda", "fetch", "xarm6_robotiq", "so100", "widowxai"].

Methods

evaluate() -> Dict[str, torch.Tensor]

Evaluates the current state of the environment and returns success metrics.

Returns:

• success (torch.Tensor): Boolean tensor indicating whether the task is successful (cube placed and robot static).
• is_obj_placed (torch.Tensor): Boolean tensor indicating whether the cube is within the goal threshold.
• is_robot_static (torch.Tensor): Boolean tensor indicating whether the robot is static (joint velocity < 0.2).
• is_grasped (torch.Tensor): Boolean tensor indicating whether the cube is currently grasped.

compute_dense_reward(obs: Any, action: torch.Tensor, info: Dict) -> torch.Tensor

Computes a dense reward signal for the current state.

The reward consists of:

• Reaching reward: Encourages the robot’s TCP (Tool Center Point) to approach the cube.
• Grasping reward: Bonus for successfully grasping the cube.
• Placement reward: Encourages moving the grasped cube to the goal position.
• Static reward: Encourages the robot to be static after placing the cube.
• Success bonus: Large reward (5.0) when the task is successful.

Args:

• obs (Any): Current observation (not used in this implementation).
• action (torch.Tensor): Action taken by the agent.
• info (Dict): Info dictionary containing is_grasped, is_obj_placed, and success keys.

Returns:

• reward (torch.Tensor): Dense reward tensor with shape (batch_size,).

compute_normalized_dense_reward(obs: Any, action: torch.Tensor, info: Dict) -> torch.Tensor

Computes a normalized dense reward signal (divided by 5.0).

Args:

• obs (Any): Current observation.
• action (torch.Tensor): Action taken by the agent.
• info (Dict): Info dictionary.

Returns:

• reward (torch.Tensor): Normalized dense reward tensor with shape (batch_size,).

_get_obs_extra(info: Dict) -> Dict[str, torch.Tensor]

Returns extra observation information.

Args:

• info (Dict): Info dictionary.

Returns:

• obs (Dict[str, torch.Tensor]): Dictionary containing:
  • is_grasped (torch.Tensor): Whether the cube is grasped.
  • tcp_pose (torch.Tensor): Tool Center Point pose.
  • goal_pos (torch.Tensor): Goal position.
  • obj_pose (torch.Tensor, optional): Cube pose (if obs_mode includes “state”).
  • tcp_to_obj_pos (torch.Tensor, optional): Vector from TCP to cube (if obs_mode includes “state”).
  • obj_to_goal_pos (torch.Tensor, optional): Vector from cube to goal (if obs_mode includes “state”).

Observation Space

The observation space depends on the obs_mode parameter. Common modes include:

• "state": State-based observations including object poses and robot state.
• "rgb": RGB camera images.
• "rgbd": RGB and depth images.

See BaseEnv documentation for more details on observation modes.

Action Space

The action space depends on the control_mode parameter and the selected robot. For the default Panda robot with control_mode="pd_joint_delta_pos", the action space is a Box(-1, 1, (8,), float32) representing normalized joint position deltas.