def _set_env_state_locomotion()

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

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

    def _set_env_state_locomotion(state: Tuple, env: gym.wrappers.TimeLimit):
        """Sets the state of the environment.

        Assumes ``state`` was generated using :func:`get_current_state`.

        Args:
            state (tuple): see :func:`get_current_state` for a description.
            env (:class:`gym.wrappers.TimeLimit`): the environment.
        """
        if _is_pybullet_gym_env(env):
            (
                filename,
                ground_ids,
                potential,
                reward,
                robot_data,
            ) = state

            env = env.env
            env.ground_ids = ground_ids
            env.potential = potential
            env.reward = reward
            PybulletEnvHandler.load_state_from_file(env._p, filename)
            for k, v in robot_data.items():
                setattr(env.robot, k, v)
        else:
            raise RuntimeError("Only pybulletgym environments supported.")