def _get_current_state_locomotion()

in mbrl/util/pybullet.py [0:0]

• 31 lines of code
• 3 McCabe index (conditional complexity)

    def _get_current_state_locomotion(env: gym.wrappers.TimeLimit) -> Tuple:
        """Returns the internal state of the environment.

        Returns a tuple with information that can be passed to :func:set_env_state` to manually
        set the environment (or a copy of it) to the same state it had when this function was
        called.

        Args:
            env (:class:`gym.wrappers.TimeLimit`): the environment.
        """
        env = env.env
        robot = env.robot
        if not isinstance(robot, (RSWalkerBase, MJWalkerBase)):
            raise RuntimeError("Invalid robot type. Expected a locomotor robot")

        filename = PybulletEnvHandler.save_state_to_file(env._p)
        ground_ids = env.ground_ids
        potential = env.potential
        reward = float(env.reward)
        robot_keys: List[Tuple[str, Callable]] = [
            ("body_rpy", tuple),
            ("body_xyz", tuple),
            ("feet_contact", np.copy),
            ("initial_z", float),
            ("joint_speeds", np.copy),
            ("joints_at_limit", int),
            ("walk_target_dist", float),
            ("walk_target_theta", float),
            ("walk_target_x", float),
            ("walk_target_y", float),
        ]

        robot_data = {}
        for k, t in robot_keys:
            robot_data[k] = t(getattr(robot, k))

        return (
            filename,
            ground_ids,
            potential,
            reward,
            robot_data,
        )