envs/block_pushing/block_pushing_multimodal.py

# coding=utf-8
# Copyright 2022 The Reach ML Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Multimodal block environments for the XArm."""

import einops
import collections
import logging
import math
from typing import Dict, List

from gym import spaces
from gym.envs import registration
from . import block_pushing
from .utils import utils_pybullet
from .utils.pose3d import Pose3d
from .utils.utils_pybullet import ObjState
from .utils.utils_pybullet import XarmState
import numpy as np
from scipy.spatial import transform
import pybullet
import pybullet_utils.bullet_client as bullet_client

import torch

# pytype: skip-file
BLOCK2_URDF_PATH = "third_party/py/envs/assets/block2.urdf"
ZONE2_URDF_PATH = "third_party/py/envs/assets/zone2.urdf"

# When resetting multiple targets, they should all be this far apart.
MIN_BLOCK_DIST = 0.1
MIN_TARGET_DIST = 0.12
# pylint: enable=line-too-long
NUM_RESET_ATTEMPTS = 1000

# Random movement of blocks
RANDOM_X_SHIFT = 0.1
RANDOM_Y_SHIFT = 0.15

logging.basicConfig(
    level="INFO",
    format="%(asctime)s [%(levelname)s] %(message)s",
    filemode="w",
)
logger = logging.getLogger()


def build_env_name(task, shared_memory, use_image_obs):
    """Construct the env name from parameters."""
    del task
    env_name = "BlockPushMultimodal"

    if use_image_obs:
        env_name = env_name + "Rgb"

    if shared_memory:
        env_name = "Shared" + env_name

    env_name = env_name + "-v0"

    return env_name


class BlockPushMultimodal(block_pushing.BlockPush):
    """2 blocks, 2 targets."""

    def __init__(
        self,
        control_frequency=10.0,
        task=block_pushing.BlockTaskVariant.PUSH,
        image_size=(224, 224),
        shared_memory=False,
        seed=None,
        goal_dist_tolerance=0.05,
    ):
        """Creates an env instance.

        Args:
          control_frequency: Control frequency for the arm. Each env step will
            advance the simulation by 1/control_frequency seconds.
          task: enum for which task, see BlockTaskVariant enum.
          image_size: Optional image size (height, width). If None, no image
            observations will be used.
          shared_memory: If True `pybullet.SHARED_MEMORY` is used to connect to
            pybullet. Useful to debug.
          seed: Optional seed for the environment.
          goal_dist_tolerance: float, how far away from the goal to terminate.
        """
        self._target_ids = None
        self._target_poses = None
        super(BlockPushMultimodal, self).__init__(
            control_frequency=control_frequency,
            task=task,
            image_size=image_size,
            shared_memory=shared_memory,
            seed=seed,
            goal_dist_tolerance=goal_dist_tolerance,
        )
        self._init_distance = [-1.0, -1.0]
        self._in_target = [[-1.0, -1.0], [-1.0, -1.0]]
        self._first_move = [-1, -1]
        self._step_num = 0
        self.moved = 0
        self.entered = 0

    @property
    def target_poses(self):
        return self._target_poses

    def get_goal_translation(self):
        """Return the translation component of the goal (2D)."""
        if self._target_poses:
            return [i.translation for i in self._target_poses]
        else:
            return None

    def _setup_pybullet_scene(self):
        self._pybullet_client = bullet_client.BulletClient(self._connection_mode)

        # Temporarily disable rendering to speed up loading URDFs.
        pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)

        self._setup_workspace_and_robot()

        self._target_ids = [
            utils_pybullet.load_urdf(self._pybullet_client, i, useFixedBase=True)
            for i in [block_pushing.ZONE_URDF_PATH, ZONE2_URDF_PATH]
        ]
        self._block_ids = []
        for i in [block_pushing.BLOCK_URDF_PATH, BLOCK2_URDF_PATH]:
            self._block_ids.append(
                utils_pybullet.load_urdf(self._pybullet_client, i, useFixedBase=False)
            )

        # Re-enable rendering.
        pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)

        self.step_simulation_to_stabilize()

    def _reset_block_poses(self, workspace_center_x):
        """Resets block poses."""

        # Helper for choosing random block position.
        def _reset_block_pose(idx, add=0.0, avoid=None):
            def _get_random_translation():
                block_x = (
                    workspace_center_x
                    + add
                    + self._rng.uniform(low=-RANDOM_X_SHIFT, high=RANDOM_X_SHIFT)
                )
                block_y = -0.2 + self._rng.uniform(
                    low=-RANDOM_Y_SHIFT, high=RANDOM_Y_SHIFT
                )
                block_translation = np.array([block_x, block_y, 0])
                return block_translation

            if avoid is None:
                block_translation = _get_random_translation()
            else:
                # Reject targets too close to `avoid`.
                for _ in range(NUM_RESET_ATTEMPTS):
                    block_translation = _get_random_translation()
                    dist = np.linalg.norm(block_translation[0] - avoid[0])
                    # print('block inner try_idx %d, dist %.3f' % (try_idx, dist))
                    if dist > MIN_BLOCK_DIST:
                        break
            block_sampled_angle = self._rng.uniform(math.pi)
            block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle])
            self._pybullet_client.resetBasePositionAndOrientation(
                self._block_ids[idx],
                block_translation.tolist(),
                block_rotation.as_quat().tolist(),
            )
            return block_translation

        # Reject targets too close to `avoid`.
        for _ in range(NUM_RESET_ATTEMPTS):
            # Reset first block.
            b0_translation = _reset_block_pose(0)
            # Reset second block away from first block.
            b1_translation = _reset_block_pose(1, avoid=b0_translation)
            dist = np.linalg.norm(b0_translation[0] - b1_translation[0])
            if dist > MIN_BLOCK_DIST:
                break
        else:
            raise ValueError("could not find matching block")
        assert dist > MIN_BLOCK_DIST

    def _reset_target_poses(self, workspace_center_x):
        """Resets target poses."""

        def _reset_target_pose(idx, add=0.0, avoid=None):
            def _get_random_translation():
                # Choose x,y randomly.
                target_x = (
                    workspace_center_x
                    + add
                    + self._rng.uniform(
                        low=-0.05 * RANDOM_X_SHIFT, high=0.05 * RANDOM_X_SHIFT
                    )
                )
                target_y = 0.2 + self._rng.uniform(
                    low=-0.05 * RANDOM_Y_SHIFT, high=0.05 * RANDOM_Y_SHIFT
                )
                target_translation = np.array([target_x, target_y, 0.020])
                return target_translation

            if avoid is None:
                target_translation = _get_random_translation()
            else:
                # Reject targets too close to `avoid`.
                for _ in range(NUM_RESET_ATTEMPTS):
                    target_translation = _get_random_translation()
                    dist = np.linalg.norm(target_translation[0] - avoid[0])
                    # print('target inner try_idx %d, dist %.3f' % (try_idx, dist))
                    if dist > MIN_TARGET_DIST:
                        break
            target_sampled_angle = math.pi + self._rng.uniform(
                low=-math.pi / 30, high=math.pi / 30
            )
            target_rotation = transform.Rotation.from_rotvec(
                [0, 0, target_sampled_angle]
            )
            self._pybullet_client.resetBasePositionAndOrientation(
                self._target_ids[idx],
                target_translation.tolist(),
                target_rotation.as_quat().tolist(),
            )
            self._target_poses[idx] = Pose3d(
                rotation=target_rotation, translation=target_translation
            )

        if self._target_poses is None:
            self._target_poses = [None for _ in range(len(self._target_ids))]

        for _ in range(NUM_RESET_ATTEMPTS):
            # Choose the first target.
            add = 0.12 * np.random.choice([-1, 1])
            # Randomly flip the location of the targets.
            _reset_target_pose(0, add=add)
            _reset_target_pose(1, add=-add, avoid=self._target_poses[0].translation)
            dist = np.linalg.norm(
                self._target_poses[0].translation[0]
                - self._target_poses[1].translation[0]
            )
            if dist > MIN_TARGET_DIST:
                break
        else:
            raise ValueError("could not find matching target")
        assert dist > MIN_TARGET_DIST

    def _reset_object_poses(self, workspace_center_x, workspace_center_y):
        # Reset block poses.
        self._reset_block_poses(workspace_center_x)

        # Reset target poses.
        self._reset_target_poses(workspace_center_x)

        self._init_distance = [-1.0, -1.0]
        self._in_target = [[-1.0, -1.0], [-1.0, -1.0]]
        self._step_num = 0

    def reset(self, reset_poses=True):
        workspace_center_x = 0.4
        workspace_center_y = 0.0
        self.moved = 0
        self.entered = 0

        if reset_poses:
            self._pybullet_client.restoreState(self._saved_state)

            rotation = transform.Rotation.from_rotvec([0, math.pi, 0])
            translation = np.array([0.3, -0.4, block_pushing.EFFECTOR_HEIGHT])
            starting_pose = Pose3d(rotation=rotation, translation=translation)
            self._set_robot_target_effector_pose(starting_pose)
            self._reset_object_poses(workspace_center_x, workspace_center_y)

        else:
            self._target_poses = [
                self._get_target_pose(idx) for idx in self._target_ids
            ]

        if reset_poses:
            self.step_simulation_to_stabilize()

        state = self._compute_state()
        self._previous_state = state
        return state

    def _get_target_pose(self, idx):
        (
            target_translation,
            target_orientation_quat,
        ) = self._pybullet_client.getBasePositionAndOrientation(idx)
        target_rotation = transform.Rotation.from_quat(target_orientation_quat)
        target_translation = np.array(target_translation)
        return Pose3d(rotation=target_rotation, translation=target_translation)

    def _compute_reach_target(self, state):
        xy_block = state["block_translation"]
        xy_target = state["target_translation"]

        xy_block_to_target = xy_target - xy_block
        xy_dir_block_to_target = (xy_block_to_target) / np.linalg.norm(
            xy_block_to_target
        )
        self.reach_target_translation = xy_block + -1 * xy_dir_block_to_target * 0.05

    def _compute_state(self):
        effector_pose = self._robot.forward_kinematics()

        def _get_block_pose(idx):
            block_position_and_orientation = (
                self._pybullet_client.getBasePositionAndOrientation(
                    self._block_ids[idx]
                )
            )
            block_pose = Pose3d(
                rotation=transform.Rotation.from_quat(
                    block_position_and_orientation[1]
                ),
                translation=block_position_and_orientation[0],
            )
            return block_pose

        block_poses = [_get_block_pose(i) for i in range(len(self._block_ids))]

        def _yaw_from_pose(pose):
            return np.array([pose.rotation.as_euler("xyz", degrees=False)[-1]])

        obs = collections.OrderedDict(
            block_translation=block_poses[0].translation[0:2],
            block_orientation=_yaw_from_pose(block_poses[0]),
            block2_translation=block_poses[1].translation[0:2],
            block2_orientation=_yaw_from_pose(block_poses[1]),
            effector_translation=effector_pose.translation[0:2],
            effector_target_translation=self._target_effector_pose.translation[0:2],
            target_translation=self._target_poses[0].translation[0:2],
            target_orientation=_yaw_from_pose(self._target_poses[0]),
            target2_translation=self._target_poses[1].translation[0:2],
            target2_orientation=_yaw_from_pose(self._target_poses[1]),
        )

        for i in range(2):
            new_distance = np.linalg.norm(
                block_poses[i].translation[0:2]
            )  # + np.linalg.norm(_yaw_from_pose(block_poses[i]))
            if self._init_distance[i] == -1:
                self._init_distance[i] = new_distance
            else:
                if self._init_distance[i] != 100:
                    if np.abs(new_distance - self._init_distance[i]) > 1e-3:
                        logger.info(f"Block {i} moved on step {self._step_num}")
                        self.moved += 1
                        self._init_distance[i] = 100

        self._step_num += 1
        return obs

    def step(self, action):
        self._step_robot_and_sim(action)

        state = self._compute_state()
        done = False
        reward = self._get_reward(state)
        if reward >= 0.5:
            # Terminate the episode if both blocks are close enough to the targets.
            done = True
        return state, reward, done, {}

    def _get_reward(self, state):
        # Reward is 1. if both blocks are inside targets, but not the same target.
        targets = ["target", "target2"]

        def _block_target_dist(block, target):
            return np.linalg.norm(
                state["%s_translation" % block] - state["%s_translation" % target]
            )

        def _closest_target(block):
            # Distances to all targets.
            dists = [_block_target_dist(block, t) for t in targets]
            # Which is closest.
            closest_target = targets[np.argmin(dists)]
            closest_dist = np.min(dists)
            # Is it in the closest target?
            in_target = closest_dist < self.goal_dist_tolerance
            return closest_target, in_target

        blocks = ["block", "block2"]

        reward = 0.0

        for t_i, t in enumerate(targets):
            for b_i, b in enumerate(blocks):
                if self._in_target[t_i][b_i] == -1:
                    dist = _block_target_dist(b, t)
                    if dist < self.goal_dist_tolerance:
                        self._in_target[t_i][b_i] = 0
                        logger.info(
                            f"Block {b_i} entered target {t_i} on step {self._step_num}"
                        )
                        self.entered += 1
                        reward += 0.49

        b0_closest_target, b0_in_target = _closest_target("block")
        b1_closest_target, b1_in_target = _closest_target("block2")
        # reward = 0.0
        if b0_in_target and b1_in_target and (b0_closest_target != b1_closest_target):
            reward = 0.51
        return reward

    def _compute_goal_distance(self, state):
        blocks = ["block", "block2"]

        def _target_block_dist(target, block):
            return np.linalg.norm(
                state["%s_translation" % block] - state["%s_translation" % target]
            )

        def _closest_block_dist(target):
            dists = [_target_block_dist(target, b) for b in blocks]
            closest_dist = np.min(dists)
            return closest_dist

        t0_closest_dist = _closest_block_dist("target")
        t1_closest_dist = _closest_block_dist("target2")
        return np.mean([t0_closest_dist, t1_closest_dist])

    @property
    def succeeded(self):
        state = self._compute_state()
        reward = self._get_reward(state)
        if reward >= 0.5:
            return True
        return False

    def _create_observation_space(self, image_size):
        pi2 = math.pi * 2

        obs_dict = collections.OrderedDict(
            block_translation=spaces.Box(low=-5, high=5, shape=(2,)),  # x,y
            block_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)),  # phi
            block2_translation=spaces.Box(low=-5, high=5, shape=(2,)),  # x,y
            block2_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)),  # phi
            effector_translation=spaces.Box(
                low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1,
                high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1,
            ),  # x,y
            effector_target_translation=spaces.Box(
                low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1,
                high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1,
            ),  # x,y
            target_translation=spaces.Box(low=-5, high=5, shape=(2,)),  # x,y
            target_orientation=spaces.Box(
                low=-pi2,
                high=pi2,
                shape=(1,),
            ),  # theta
            target2_translation=spaces.Box(low=-5, high=5, shape=(2,)),  # x,y
            target2_orientation=spaces.Box(
                low=-pi2,
                high=pi2,
                shape=(1,),
            ),  # theta
        )
        if image_size is not None:
            obs_dict["rgb"] = spaces.Box(
                low=0, high=255, shape=(image_size[0], image_size[1], 3), dtype=np.uint8
            )
        return spaces.Dict(obs_dict)

    def get_pybullet_state(self):
        """Save pybullet state of the scene.

        Returns:
          dict containing 'robots', 'robot_end_effectors', 'targets', 'objects',
            each containing a list of ObjState.
        """
        state: Dict[str, List[ObjState]] = {}

        state["robots"] = [
            XarmState.get_bullet_state(
                self._pybullet_client,
                self.robot.xarm,
                target_effector_pose=self._target_effector_pose,
                goal_translation=None,
            )
        ]

        state["robot_end_effectors"] = []
        if self.robot.end_effector:
            state["robot_end_effectors"].append(
                ObjState.get_bullet_state(
                    self._pybullet_client, self.robot.end_effector
                )
            )

        state["targets"] = []
        if self._target_ids:
            for target_id in self._target_ids:
                state["targets"].append(
                    ObjState.get_bullet_state(self._pybullet_client, target_id)
                )

        state["objects"] = []
        for obj_id in self.get_obj_ids():
            state["objects"].append(
                ObjState.get_bullet_state(self._pybullet_client, obj_id)
            )

        return state

    def set_pybullet_state(self, state):
        """Restore pyullet state.

        WARNING: py_environment wrapper assumes environments aren't reset in their
        constructor and will often reset the environment unintentionally. It is
        always recommeneded that you call env.reset on the tfagents wrapper before
        playback (replaying pybullet_state).

        Args:
          state: dict containing 'robots', 'robot_end_effectors', 'targets',
            'objects', each containing a list of ObjState.
        """

        assert isinstance(state["robots"][0], XarmState)
        xarm_state: XarmState = state["robots"][0]
        xarm_state.set_bullet_state(self._pybullet_client, self.robot.xarm)
        self._set_robot_target_effector_pose(xarm_state.target_effector_pose)

        def _set_state_safe(obj_state, obj_id):
            if obj_state is not None:
                assert obj_id is not None, "Cannot set state for missing object."
                obj_state.set_bullet_state(self._pybullet_client, obj_id)
            else:
                assert obj_id is None, f"No state found for obj_id {obj_id}"

        robot_end_effectors = state["robot_end_effectors"]
        _set_state_safe(
            None if not robot_end_effectors else robot_end_effectors[0],
            self.robot.end_effector,
        )

        for target_state, target_id in zip(state["targets"], self._target_ids):
            _set_state_safe(target_state, target_id)

        obj_ids = self.get_obj_ids()
        assert len(state["objects"]) == len(obj_ids), "State length mismatch"
        for obj_state, obj_id in zip(state["objects"], obj_ids):
            _set_state_safe(obj_state, obj_id)

        self.reset(reset_poses=False)


class BlockPushMultimodalMultiview(BlockPushMultimodal):
    def __init__(self, id=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.observation_space = spaces.Box(
            low=0, high=1, shape=(2, 3, self._image_size[0], self._image_size[1])
        )
        self._step = 0

    def _get_obs(self):
        # render to VCHW shape
        view0 = self._render_camera(self._image_size, view=0)
        view1 = self._render_camera(self._image_size, view=1)
        obs = np.stack([view0, view1], axis=0)  # VHWC
        return einops.rearrange(obs, "V H W C -> V C H W")

    def step(self, action):
        action = action * 0.03
        self._step_robot_and_sim(action)

        state = self._compute_state()
        reward = self._get_reward(state)
        # Terminate the episode if both blocks are close enough to the targets.
        obs = self._get_obs()
        image = einops.rearrange(obs, "V C H W -> H (V W) C")
        obs = obs / 255.0
        self._step += 1
        done = (reward >= 0.5) or (self._step >= 300)

        return (
            obs,
            reward,
            done,
            {
                "state": state,
                "image": image,
                "all_completions_ids": [],
                "moved": self.moved,
                "entered": self.entered,
            },
        )

    def reset(self, reset_poses=True, *args, **kwargs):
        print("resetting env")
        self._step = 0
        state = super().reset(reset_poses=reset_poses)
        obs = self._get_obs()
        obs = obs / 255.0
        return obs

    def set_state(self, state: collections.OrderedDict):
        robot_t = np.array(
            [*state["effector_translation"], block_pushing.EFFECTOR_HEIGHT]
        )
        robot_r = transform.Rotation.from_rotvec([0, np.pi, 0])
        robot_pose = Pose3d(rotation=robot_r, translation=robot_t)
        self._set_robot_target_effector_pose(robot_pose)
        self.step_simulation_to_stabilize()

        block_t = [*state["block_translation"], 0]
        block_r = transform.Rotation.from_rotvec([0, 0, state["block_orientation"]])
        self._pybullet_client.resetBasePositionAndOrientation(
            self._block_ids[0],
            block_t,
            block_r.as_quat().tolist(),
        )

        block2_t = [*state["block2_translation"], 0]
        block2_r = transform.Rotation.from_rotvec([0, 0, state["block2_orientation"]])
        self._pybullet_client.resetBasePositionAndOrientation(
            self._block_ids[1],
            block2_t,
            block2_r.as_quat().tolist(),
        )

        target_t = [*state["target_translation"], 0.02]
        target_r = transform.Rotation.from_rotvec([0, 0, state["target_orientation"]])
        self._pybullet_client.resetBasePositionAndOrientation(
            self._target_ids[0],
            target_t,
            target_r.as_quat().tolist(),
        )

        target2_t = [*state["target2_translation"], 0.02]
        target2_r = transform.Rotation.from_rotvec([0, 0, state["target2_orientation"]])
        self._pybullet_client.resetBasePositionAndOrientation(
            self._target_ids[1],
            target2_t,
            target2_r.as_quat().tolist(),
        )

        self.step_simulation_to_stabilize()


class BlockPushHorizontalMultimodal(BlockPushMultimodal):
    def _reset_object_poses(self, workspace_center_x, workspace_center_y):
        # Reset block poses.
        self._reset_block_poses(workspace_center_y)

        # Reset target poses.
        self._reset_target_poses(workspace_center_y)

    def _reset_block_poses(self, workspace_center_y):
        """Resets block poses."""

        # Helper for choosing random block position.
        def _reset_block_pose(idx, add=0.0, avoid=None):
            def _get_random_translation():
                block_x = 0.35 + 0.5 * self._rng.uniform(
                    low=-RANDOM_X_SHIFT, high=RANDOM_X_SHIFT
                )
                block_y = (
                    workspace_center_y
                    + add
                    + 0.5 * self._rng.uniform(low=-RANDOM_Y_SHIFT, high=RANDOM_Y_SHIFT)
                )
                block_translation = np.array([block_x, block_y, 0])
                return block_translation

            if avoid is None:
                block_translation = _get_random_translation()
            else:
                # Reject targets too close to `avoid`.
                for _ in range(NUM_RESET_ATTEMPTS):
                    block_translation = _get_random_translation()
                    dist = np.linalg.norm(block_translation[0] - avoid[0])
                    # print('block inner try_idx %d, dist %.3f' % (try_idx, dist))
                    if dist > MIN_BLOCK_DIST:
                        break
            block_sampled_angle = self._rng.uniform(math.pi)
            block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle])
            self._pybullet_client.resetBasePositionAndOrientation(
                self._block_ids[idx],
                block_translation.tolist(),
                block_rotation.as_quat().tolist(),
            )
            return block_translation

        # Reject targets too close to `avoid`.
        for _ in range(NUM_RESET_ATTEMPTS):
            # Reset first block.
            add = 0.2 * np.random.choice([-1, 1])
            b0_translation = _reset_block_pose(0, add=add)
            # Reset second block away from first block.
            b1_translation = _reset_block_pose(1, add=-add, avoid=b0_translation)
            dist = np.linalg.norm(b0_translation[0] - b1_translation[0])
            if dist > MIN_BLOCK_DIST:
                break
        else:
            raise ValueError("could not find matching block")
        assert dist > MIN_BLOCK_DIST

    def _reset_target_poses(self, workspace_center_y):
        """Resets target poses."""

        def _reset_target_pose(idx, add=0.0, avoid=None):
            def _get_random_translation():
                # Choose x,y randomly.
                target_x = 0.5 + self._rng.uniform(
                    low=-0.05 * RANDOM_X_SHIFT, high=0.05 * RANDOM_X_SHIFT
                )
                target_y = (
                    workspace_center_y
                    + add
                    + self._rng.uniform(
                        low=-0.05 * RANDOM_Y_SHIFT, high=0.05 * RANDOM_Y_SHIFT
                    )
                )
                target_translation = np.array([target_x, target_y, 0.020])
                return target_translation

            if avoid is None:
                target_translation = _get_random_translation()
            else:
                # Reject targets too close to `avoid`.
                for _ in range(NUM_RESET_ATTEMPTS):
                    target_translation = _get_random_translation()
                    dist = np.linalg.norm(target_translation[0] - avoid[0])
                    # print('target inner try_idx %d, dist %.3f' % (try_idx, dist))
                    if dist > MIN_TARGET_DIST:
                        break
            target_sampled_angle = math.pi + self._rng.uniform(
                low=-math.pi / 30, high=math.pi / 30
            )
            target_rotation = transform.Rotation.from_rotvec(
                [0, 0, target_sampled_angle]
            )
            self._pybullet_client.resetBasePositionAndOrientation(
                self._target_ids[idx],
                target_translation.tolist(),
                target_rotation.as_quat().tolist(),
            )
            self._target_poses[idx] = Pose3d(
                rotation=target_rotation, translation=target_translation
            )

        if self._target_poses is None:
            self._target_poses = [None for _ in range(len(self._target_ids))]

        for _ in range(NUM_RESET_ATTEMPTS):
            # Choose the first target.
            add = 0.2 * np.random.choice([-1, 1])
            # Randomly flip the location of the targets.
            _reset_target_pose(0, add=add)
            _reset_target_pose(1, add=-add, avoid=self._target_poses[0].translation)
            dist = np.linalg.norm(
                self._target_poses[0].translation[0]
                - self._target_poses[1].translation[0]
            )
            break
            # if dist > MIN_TARGET_DIST:
            #     break
        else:
            raise ValueError("could not find matching target")
        # assert dist > MIN_TARGET_DIST


if "BlockPushMultimodal-v0" in registration.registry.env_specs:
    del registration.registry["BlockPushMultimodal-v0"]

registration.register(
    id="BlockPushMultimodal-v0", entry_point=BlockPushMultimodal, max_episode_steps=350
)

registration.register(
    id="BlockPushMultimodalFlipped-v0",
    entry_point=BlockPushHorizontalMultimodal,
    max_episode_steps=25,
)

registration.register(
    id="SharedBlockPushMultimodal-v0",
    entry_point=BlockPushMultimodal,
    kwargs=dict(shared_memory=True),
    max_episode_steps=350,
)
registration.register(
    id="BlockPushMultimodalRgb-v0",
    entry_point=BlockPushMultimodal,
    max_episode_steps=350,
    kwargs=dict(image_size=(block_pushing.IMAGE_HEIGHT, block_pushing.IMAGE_WIDTH)),
)
registration.register(
    id="BlockPushMultimodalMultiview-v0",
    entry_point=BlockPushMultimodalMultiview,
    max_episode_steps=350,
    kwargs=dict(image_size=(block_pushing.IMAGE_HEIGHT, block_pushing.IMAGE_WIDTH)),
)