diff --git a/CMakeLists.txt b/CMakeLists.txt index 91fbe9530..b1d5b9449 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.10) # Set the build version -set(BUILD_VERSION 1.4.11) +set(BUILD_VERSION 1.4.12) # Add definition of Jiminy version for C++ headers add_definitions("-DJIMINY_VERSION=\"${BUILD_VERSION}\"") diff --git a/core/src/control/AbstractController.cc b/core/src/control/AbstractController.cc index 3fc41fefb..49f14434d 100644 --- a/core/src/control/AbstractController.cc +++ b/core/src/control/AbstractController.cc @@ -33,8 +33,11 @@ namespace jiminy return hresult_t::ERROR_INIT_FAILED; } + // Backup robot robot_ = robot; - sensorsData_ = robot_->getSensorsData(); + + // Reset the controller completely + reset(true); try { diff --git a/core/src/engine/Engine.cc b/core/src/engine/Engine.cc index 0ea03d285..e315331fa 100644 --- a/core/src/engine/Engine.cc +++ b/core/src/engine/Engine.cc @@ -23,6 +23,19 @@ namespace jiminy { hresult_t returnCode = hresult_t::SUCCESS; + // Make sure the simulation is properly stopped + if (isSimulationRunning_) + { + stop(); + } + + // Remove the existing system if already initialized + if(isInitialized_) + { + removeSystem(""); // It cannot fail at this point + isInitialized_ = false; + } + /* Add the system without associated name, since it is irrelevant for a single robot engine. */ returnCode = addSystem("", std::move(robot), diff --git a/core/src/engine/EngineMultiRobot.cc b/core/src/engine/EngineMultiRobot.cc index ee4f1a9b1..cb62d32b0 100644 --- a/core/src/engine/EngineMultiRobot.cc +++ b/core/src/engine/EngineMultiRobot.cc @@ -67,6 +67,13 @@ namespace jiminy std::shared_ptr controller, callbackFunctor_t callbackFct) { + // Make sure that no simulation is running + if (isSimulationRunning_) + { + PRINT_ERROR("A simulation is already running. Stop it before adding a new system.") + return hresult_t::ERROR_GENERIC; + } + if (!robot->getIsInitialized()) { PRINT_ERROR("Robot not initialized.") @@ -119,8 +126,18 @@ namespace jiminy { hresult_t returnCode = hresult_t::SUCCESS; - // Remove every coupling forces involving the system - returnCode = removeCouplingForces(systemName); + // Make sure that no simulation is running + if (isSimulationRunning_) + { + PRINT_ERROR("A simulation is already running. Stop it before removing a system.") + returnCode = hresult_t::ERROR_GENERIC; + } + + if (returnCode == hresult_t::SUCCESS) + { + // Remove every coupling forces involving the system + returnCode = removeCouplingForces(systemName); + } if (returnCode == hresult_t::SUCCESS) { @@ -195,8 +212,18 @@ namespace jiminy { hresult_t returnCode = hresult_t::SUCCESS; + // Make sure that no simulation is running + if (isSimulationRunning_) + { + PRINT_ERROR("A simulation is already running. Stop it before adding coupling forces.") + returnCode = hresult_t::ERROR_GENERIC; + } + int32_t systemIdx1; - returnCode = getSystemIdx(systemName1, systemIdx1); + if (returnCode == hresult_t::SUCCESS) + { + returnCode = getSystemIdx(systemName1, systemIdx1); + } int32_t systemIdx2; if (returnCode == hresult_t::SUCCESS) @@ -239,8 +266,18 @@ namespace jiminy { hresult_t returnCode = hresult_t::SUCCESS; + // Make sure that no simulation is running + if (isSimulationRunning_) + { + PRINT_ERROR("A simulation is already running. Stop it before removing coupling forces.") + returnCode = hresult_t::ERROR_GENERIC; + } + systemHolder_t * system1; - returnCode = getSystem(systemName1, system1); + if (returnCode == hresult_t::SUCCESS) + { + returnCode = getSystem(systemName1, system1); + } if (returnCode == hresult_t::SUCCESS) { @@ -268,8 +305,18 @@ namespace jiminy { hresult_t returnCode = hresult_t::SUCCESS; + // Make sure that no simulation is running + if (isSimulationRunning_) + { + PRINT_ERROR("A simulation is already running. Stop it before removing coupling forces.") + returnCode = hresult_t::ERROR_GENERIC; + } + systemHolder_t * system; - returnCode = getSystem(systemName, system); + if (returnCode == hresult_t::SUCCESS) + { + returnCode = getSystem(systemName, system); + } if (returnCode == hresult_t::SUCCESS) { @@ -580,9 +627,6 @@ namespace jiminy // Reset the robot, controller, engine, and registered impulse forces if requested reset(resetRandomNumbers, resetDynamicForceRegister); - // At this point, consider that the simulation is running - isSimulationRunning_ = true; - auto systemIt = systems_.begin(); auto systemDataIt = systemsDataHolder_.begin(); for ( ; systemIt != systems_.end(); ++systemIt, ++systemDataIt) @@ -807,6 +851,14 @@ namespace jiminy } } + // At this point, consider that the simulation is running + isSimulationRunning_ = true; + + if (returnCode != hresult_t::SUCCESS) + { + stop(); + } + return returnCode; } @@ -2129,7 +2181,7 @@ namespace jiminy // Convert contact force from the global frame to the local frame to store it in contactForces_ pinocchio::SE3 const & transformContactInJoint = system.robot->pncModel_.frames[frameIdx].placement; - system.robot->contactForces_[i] = transformContactInJoint.act(fextLocal); + system.robot->contactForces_[i] = transformContactInJoint.actInv(fextLocal); } // Compute the force at collision bodies diff --git a/python/gym_jiminy/common/gym_jiminy/common/block_bases.py b/python/gym_jiminy/common/gym_jiminy/common/block_bases.py index 9f036b0ee..f0ee0feb5 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/block_bases.py +++ b/python/gym_jiminy/common/gym_jiminy/common/block_bases.py @@ -208,7 +208,9 @@ def _refresh_observation_space(self) -> None: it, then using `BaseObserverBlock` or `BlockInterface` directly is probably the way to go. """ + # Assertion(s) for type checker assert self.env is not None + self.observation_space = self.env.action_space def reset(self, env: Union[gym.Wrapper, BaseJiminyEnv]) -> None: @@ -219,10 +221,11 @@ def reset(self, env: Union[gym.Wrapper, BaseJiminyEnv]) -> None: :param env: Environment to control, eventually already wrapped. """ + # Call base implementation super().reset(env) # Assertion(s) for type checker - assert self.env is not None and self.env.control_dt is not None + assert self.env is not None self.control_dt = self.env.control_dt * self.update_ratio @@ -334,7 +337,9 @@ def _refresh_action_space(self) -> None: it, then using `BaseControllerBlock` or `BlockInterface` directly is probably the way to go. """ + # Assertion(s) for type checker assert self.env is not None + self.action_space = self.env.observation_space def reset(self, env: Union[gym.Wrapper, BaseJiminyEnv]) -> None: @@ -345,10 +350,11 @@ def reset(self, env: Union[gym.Wrapper, BaseJiminyEnv]) -> None: :param env: Environment to observe, eventually already wrapped. """ + # Call base implementation super().reset(env) # Assertion(s) for type checker - assert self.env is not None and self.env.observe_dt is not None + assert self.env is not None self.observe_dt = self.env.observe_dt * self.update_ratio diff --git a/python/gym_jiminy/common/gym_jiminy/common/control_impl.py b/python/gym_jiminy/common/gym_jiminy/common/control_impl.py index 44796d204..6207235a5 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/control_impl.py +++ b/python/gym_jiminy/common/gym_jiminy/common/control_impl.py @@ -77,9 +77,6 @@ def _setup(self) -> None: It updates the mapping from motors to encoders indices. """ - # Assertion(s) for type checker - assert self.robot is not None and self.system_state is not None - # Refresh the mapping between the motors and encoders encoder_joints = [] for name in self.robot.sensors_names[encoder.type]: diff --git a/python/gym_jiminy/common/gym_jiminy/common/env_bases.py b/python/gym_jiminy/common/gym_jiminy/common/env_bases.py index f1ee7a599..390df9e89 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/env_bases.py +++ b/python/gym_jiminy/common/gym_jiminy/common/env_bases.py @@ -20,11 +20,11 @@ from jiminy_py.dynamics import compute_freeflyer_state_from_fixed_body from jiminy_py.simulator import Simulator from jiminy_py.viewer import sleep, play_logfiles -from jiminy_py.controller import BaseJiminyController +from jiminy_py.controller import BaseJiminyController, ControllerHandleType from pinocchio import neutral -from .utils import _clamp, zeros, SpaceDictRecursive +from .utils import _clamp, zeros, fill, SpaceDictRecursive from .generic_bases import ObserveAndControlInterface from .play import loop_interactive @@ -62,19 +62,15 @@ class BaseJiminyEnv(gym.Env, ObserveAndControlInterface): 'render.modes': ['human', 'rgb_array'], } - simulator: Optional[Simulator] - def __init__(self, - simulator: Optional[Simulator], + simulator: Simulator, step_dt: float, enforce_bounded: Optional[bool] = False, debug: bool = False, **kwargs: Any) -> None: r""" :param simulator: Jiminy Python simulator used for physics - computations. Can be `None` if `_setup` has been - overwritten such that 'self.simulator' is a valid and - completely initialized engine. + computations. :param step_dt: Simulation timestep for learning. Note that it is independent from the controller and observation update periods. The latter are configured via @@ -94,7 +90,7 @@ def __init__(self, # pylint: disable=unused-argument # Initialize the interfaces through multiple inheritance - super().__init__() + super().__init__(**kwargs) # Backup some user arguments self.simulator = simulator @@ -108,40 +104,43 @@ def __init__(self, self._log_data: Optional[Dict[str, np.ndarray]] = None self.log_path: Optional[str] = None - # Current observation and action of the robot - self._state: Optional[Tuple[np.ndarray, np.ndarray]] = None - self._sensors_data: Optional[Dict[str, np.ndarray]] = None + # Initialize sensors data shared memory + self._sensors_data = OrderedDict(self.robot.sensors_data) # Information about the learning process self._info: Dict[str, Any] = {} # Number of simulation steps performed self.num_steps = -1 - self.max_steps: Optional[int] = None + self.max_steps = int( + self.simulator.simulation_duration_max // self.step_dt) self._num_steps_beyond_done: Optional[int] = None - # Set the seed of the simulation and reset the environment. - # Note that `reset` is called to make sure the environment is - # completely configured at initialization, which is required by many - # external libraries since it is an implicit required on the OpenAI - # Gym standard interface. + # Initialize the seed of the environment self.seed() - self.reset() + + # Refresh the observation and action spaces + self._refresh_observation_space() + self._refresh_action_space() + + # Assertion(s) for type checker + assert (self.observation_space is not None and + self.action_space is not None) + + # Initialize some internal buffers + self._action = zeros(self.action_space) + self._state = (np.zeros(self.robot.nq), np.zeros(self.robot.nv)) + self._observation = zeros(self.observation_space) @property def robot(self) -> jiminy.Robot: - """ TODO: Write documentation. + """ Get robot. """ - if self.simulator is None: - raise RuntimeError("Backend simulator undefined.") return self.simulator.robot def _get_time_space(self) -> gym.Space: """Get time space. """ - # Assertion(s) for type checker - assert self.simulator is not None - return gym.spaces.Box( low=0.0, high=self.simulator.simulation_duration_max, shape=(1,), dtype=np.float64) @@ -162,9 +161,6 @@ def _get_state_space(self, value 'simulator.use_theoretical_model'. Optional: `None` by default. """ - # Assertion(s) for type checker - assert self.simulator is not None and self.robot is not None - # Handling of default argument if use_theoretical_model is None: use_theoretical_model = self.simulator.use_theoretical_model @@ -379,99 +375,65 @@ def _refresh_action_space(self) -> None: high=effort_limit[motors_velocity_idx], dtype=np.float64) - def set_state(self, qpos: np.ndarray, qvel: np.ndarray) -> None: - """Reset the simulation and specify the initial state of the robot. + def reset(self, + controller_hook: Optional[ + Callable[[], Optional[ControllerHandleType]]] = None + ) -> SpaceDictRecursive: + """Reset the environment. + In practice, it resets the backend simulator and set the initial state + of the robot. The initial state is obtained by calling '_sample_state'. This method is also in charge of setting the initial action (at the beginning) and observation (at the end). .. warning:: It starts the simulation immediately. As a result, it is not - possible to changed the robot (included options), nor to register - log variable. Yet, it is possible to do it passing a custom - 'controller_hook' to `reset` method. - - :param qpos: Configuration of the robot. - :param qvel: Velocity vector of the robot. - """ - # Assertion(s) for type checker - assert self.simulator is not None - - # Reset the simulator - self.simulator.reset() - - # Set default action. It will be used for the initial step. - self._action = zeros(self.action_space) - - # Start the engine, in order to initialize the sensors data - hresult = self.simulator.start( - qpos, qvel, self.simulator.use_theoretical_model) - if hresult != jiminy.hresult_t.SUCCESS: - raise RuntimeError( - "Failed to start the simulation. Probably because the " - "initial state is invalid.") - - # Initialize some internal buffers - self.num_steps = 0 - self.max_steps = int( - self.simulator.simulation_duration_max // self.step_dt) - self._num_steps_beyond_done = None - self._log_data = None - - # Create a new log file - if self.debug: - fd, self.log_path = tempfile.mkstemp(prefix="log_", suffix=".data") - os.close(fd) - - # - self._state = (qpos, qvel) - - # Update the observation - self._observation = self.compute_observation() - - def reset(self, - controller_hook: Optional[Callable[[], None]] = None - ) -> SpaceDictRecursive: - """Reset the environment. - - The initial state is obtained by calling '_sample_state'. + possible to change the robot (included options), nor to register + log variable. The only way to do so is via 'controller_hook'. :param controller_hook: Custom controller hook. It will be executed right after initialization of the environment, and just before actually starting the - simulation. It is useful to override partially + simulation. It is a callable that optionally + returns a controller handle. If so, it will be + used to initialize the low-level jiminy + controller. It is useful to override partially the configuration of the low-level engine, set a custom low-level controller handle, or to - register custom variables to the telemetry. - `None` to disable. + register custom variables to the telemetry. Set + to `None` if unused. Optional: Disable by default. :returns: Initial observation of the episode. """ # pylint: disable=arguments-differ - # Define default controller hook - def register() -> None: - nonlocal self - - # Assertion(s) for type checker - assert self.simulator is not None - - self.simulator.controller.set_controller_handle( - self._send_command) + # Assertion(s) for type checker + assert self.observation_space is not None - # Stop simulator if still running - if self.simulator is not None: - self.simulator.stop() + # Reset the simulator + self.simulator.reset() # Make sure the environment is properly setup self._setup() - # Initialize sensors data shared memory + # Re-initialize sensors data shared memory. + # It must be done because the robot may have changed. self._sensors_data = OrderedDict(self.robot.sensors_data) - # Set the seed of the simulator - self.simulator.seed(self._seed) + # Set default action. + # It will be used for the initial step. + fill(self._action, 0.0) + + # Reset some internal buffers + self.num_steps = 0 + self._num_steps_beyond_done = None + self._log_data = None + + # Create a new log file + if self.debug: + fd, self.log_path = tempfile.mkstemp(prefix="log_", suffix=".data") + os.close(fd) # Extract the controller and observer update period. # There is no actual observer by default, apart from the robot's state @@ -481,35 +443,30 @@ def register() -> None: engine_options = self.simulator.engine.get_options() self.control_dt = \ float(engine_options['stepper']['controllerUpdatePeriod']) - self.observer_dt = \ + self.observe_dt = \ float(engine_options['stepper']['sensorsUpdatePeriod']) - # Refresh the observation and action spaces - self._refresh_observation_space() - self._refresh_action_space() - - # Assertion(s) for type checker - assert self.observation_space is not None - - # Initialize some internal buffers - self._state = (np.zeros(self.robot.nq), np.zeros(self.robot.nv)) - self._observation = zeros(self.observation_space) - self._action = zeros(self.action_space) - # Enforce the low-level controller. - # Note that `BaseJiminyController` is used by default instead of - # `jiminy.ControllerFunctor`. Although it is less efficient because - # it adds an extra layer of indirection, it makes it possible to update - # the controller handle without instantiating a new controller. + # The backend robot may have changed, for example if it is randomly + # generated based on different URDF files. As a result, it is necessary + # to instantiate a new low-level controller. + # Note that `BaseJiminyController` is used by default in place of + # `jiminy.ControllerFunctor`. Although it is less efficient because it + # adds an extra layer of indirection, it makes it possible to update + # the controller handle without instantiating a new controller, which + # is necessary to allow registering telemetry variables before knowing + # the controller handle in advance. controller = BaseJiminyController() controller.initialize(self.robot) self.simulator.set_controller(controller) - # Run controller hook - if controller_hook is None: - register() - else: - controller_hook() + # Run controller hook and set the controller handle + controller_handle = None + if controller_hook is not None: + controller_handle = controller_hook() + if controller_handle is None: + controller_handle = self._send_command + self.simulator.controller.set_controller_handle(controller_handle) # Sample the initial state and reset the low-level engine qpos, qvel = self._sample_state() @@ -519,7 +476,17 @@ def register() -> None: "The initial state provided by `_sample_state` is " "inconsistent with the dimension or types of joints of the " "model.") - self.set_state(qpos, qvel) + + # Start the engine, in order to initialize the sensors data + hresult = self.simulator.start( + qpos, qvel, self.simulator.use_theoretical_model) + if hresult != jiminy.hresult_t.SUCCESS: + raise RuntimeError( + "Failed to start the simulation. Probably because the " + "initial state is invalid.") + + # Update the observation + self._observation = self.compute_observation() # Make sure the state is valid, otherwise there `compute_observation` # and s`_refresh_observation_space` are probably inconsistent. @@ -562,14 +529,16 @@ def seed(self, seed: Optional[int] = None) -> Sequence[np.uint32]: self._seed = np.uint32( seeding._int_list_from_bigint(seeding.hash_seed(self._seed))[0]) + # Reset the seed of Jiminy Engine + self.simulator.seed(self._seed) + return [self._seed] def close(self) -> None: """Terminate the Python Jiminy engine. Mostly defined for compatibility with Gym OpenAI. """ - if self.simulator is not None: - self.simulator.close() + self.simulator.close() def step(self, action: Optional[np.ndarray] = None @@ -581,11 +550,6 @@ def step(self, :returns: Next observation, reward, status of the episode (done or not), and a dictionary of extra information """ - # Assertion(s) for type checker - assert (self.simulator is not None and - self.max_steps is not None and - isinstance(self._action, np.ndarray)) - # Make sure a simulation is already running if not self.simulator.is_simulation_running: raise RuntimeError( @@ -656,9 +620,10 @@ def step(self, def get_log(self) -> Tuple[Dict[str, np.ndarray], Dict[str, str]]: """Get log of recorded variable since the beginning of the episode. """ - # Assertion(s) for type checker - assert self.simulator is not None - + if not self.simulator.is_simulation_running: + raise RuntimeError( + "No simulation running. Please call `reset` at least one " + "before getting log.") return self.simulator.get_log() def render(self, @@ -678,9 +643,6 @@ def render(self, :returns: RGB array if 'mode' is 'rgb_array', None otherwise. """ - # Assertion(s) for type checker - assert self.simulator is not None - if mode == 'human': return_rgb_array = False elif mode == 'rgb_array': @@ -696,9 +658,6 @@ def replay(self, **kwargs: Any) -> None: :param kwargs: Extra keyword arguments for `play_logfiles` delegation. """ - # Assertion(s) for type checker - assert self.simulator is not None and self.robot is not None - if self._log_data is not None: log_data = self._log_data else: @@ -743,16 +702,8 @@ def _setup(self) -> None: By default, it enforces some options of the engine. .. note:: - This method is called internally by `reset` method at the very - beginning. This method can be overwritten to postpone the engine - and robot creation at `reset`. One have to delegate the creation - and initialization of the engine to this method, so that it - alleviates the requirement to specify a valid the engine during the - instantiation of the environment. + This method is called internally by `reset` methods. """ - # Assertion(s) for type checker - assert self.simulator is not None and self.robot is not None - # Extract some proxies robot_options = self.robot.get_options() engine_options = self.simulator.engine.get_options() @@ -760,7 +711,7 @@ def _setup(self) -> None: # Disable part of the telemetry in non debug mode, to speed up the # simulation. Only the required data for log replay are enabled. It is # up to the user to overload this method if logging more data is - # necessary for terminal reward computation. + # necessary for computating the terminal reward. for field in robot_options["telemetry"].keys(): robot_options["telemetry"][field] = self.debug for field in engine_options["telemetry"].keys(): @@ -820,9 +771,6 @@ def _neutral(self) -> np.ndarray: initial state. It can be overloaded to ensure static stability of the configuration. """ - # Assertion(s) for type checker - assert self.simulator is not None and self.robot is not None - # Get the neutral configuration of the actual model qpos = neutral(self.robot.pinocchio_model) @@ -852,9 +800,6 @@ def _sample_state(self) -> Tuple[np.ndarray, np.ndarray]: This method is called internally by `reset` to generate the initial state. It can be overloaded to act as a random state generator. """ - # Assertion(s) for type checker - assert self.simulator is not None and self.robot is not None - # Get the neutral configuration qpos = self._neutral() @@ -876,9 +821,6 @@ def compute_observation(self # type: ignore[override] """ # pylint: disable=arguments-differ - # Assertion(s) for type checker - assert self.simulator is not None - # Update some internal buffers if self.simulator.is_simulation_running: self._state = self.simulator.state @@ -972,8 +914,12 @@ def __init__(self, simulator: Optional[Simulator], step_dt: float, debug: bool = False) -> None: + # Initialize base class super().__init__(simulator, step_dt, debug) + # Define some internal buffers + self._desired_goal: Optional[np.ndarray] = None + def _refresh_observation_space(self) -> None: # Assertion(s) for type checker assert isinstance(self._desired_goal, np.ndarray), ( @@ -994,6 +940,9 @@ def _refresh_observation_space(self) -> None: def compute_observation(self # type: ignore[override] ) -> SpaceDictRecursive: + # Assertion(s) for type checker + assert self._desired_goal is not None + return OrderedDict( observation=super().compute_observation(), achieved_goal=self._get_achieved_goal(), diff --git a/python/gym_jiminy/common/gym_jiminy/common/env_locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/env_locomotion.py index 4831daf8f..73ff43813 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/env_locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/env_locomotion.py @@ -70,8 +70,6 @@ class WalkerJiminyEnv(BaseJiminyEnv): 'render.modes': ['human'], } - simulator: Optional[Simulator] - def __init__(self, urdf_path: str, hardware_path: Optional[str] = None, @@ -151,40 +149,39 @@ def __init__(self, self._forces_profile: Sequence[ForceProfileType] = [] self._f_xy_profile_spline: Optional[ nb.core.dispatcher.Dispatcher] = None - self._power_consumption_max: Optional[float] = None - self._height_neutral: Optional[float] = None + self._power_consumption_max = 0.0 + self._height_neutral = 0.0 # Configure and initialize the learning environment - super().__init__(None, step_dt, enforce_bounded, debug, **kwargs) + simulator = Simulator.build( + self.urdf_path, self.hardware_path, self.mesh_path, + has_freeflyer=True, use_theoretical_model=False, + config_path=self.config_path, + avoid_instable_collisions=self.avoid_instable_collisions, + debug=debug) + super().__init__(simulator, step_dt, enforce_bounded, debug, **kwargs) def _setup(self) -> None: """Configure the environment. It is doing the following steps, successively: - - creates a low-level engine is necessary, - updates some proxies that will be used for computing the reward and termination condition, - enforce some options of the low-level robot and engine, - randomize the environment according to 'std_ratio'. - .. note:: + .. note:: TODO WRONG This method is called internally by `reset` method at the very - beginning. This method can be overwritten to implement new - contributions to the environment stochasticity, or to create - custom low-level robot if the model must be different for each - learning eposide for some reason. + beginning. One must overide it to implement new contributions to + the environment stochasticity, or to create custom low-level robot + if the model must be different for each learning episode. """ + # Call the base implementation + super()._setup() + # Check that a valid engine is available, and if not, create one - if self.simulator is None: - self.simulator = Simulator.build( - self.urdf_path, self.hardware_path, self.mesh_path, - has_freeflyer=True, use_theoretical_model=False, - config_path=self.config_path, - avoid_instable_collisions=self.avoid_instable_collisions, - debug=self.debug) - else: - self.simulator.remove_forces() + self.simulator.remove_forces() if not self.robot.has_freeflyer: raise RuntimeError( @@ -381,11 +378,6 @@ def is_done(self) -> bool: # type: ignore[override] """ # pylint: disable=arguments-differ - # Assertion(s) for type checker - assert (self.simulator is not None and - self._state is not None and - self._height_neutral is not None) - if self._state[0][2] < self._height_neutral * 0.75: return True if self.simulator.stepper_state.t >= self.simu_duration_max: @@ -407,10 +399,6 @@ def compute_reward(self, # type: ignore[override] """ # pylint: disable=arguments-differ - # Assertion(s) for type checker - assert (self.simulator is not None and - self._power_consumption_max is not None) - reward_dict = info.setdefault('reward', {}) # Define some proxies diff --git a/python/gym_jiminy/common/gym_jiminy/common/generic_bases.py b/python/gym_jiminy/common/gym_jiminy/common/generic_bases.py index 774ae1298..b581f59df 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/generic_bases.py +++ b/python/gym_jiminy/common/gym_jiminy/common/generic_bases.py @@ -13,7 +13,7 @@ class ControlInterface: """Controller interface for both controllers and environments. """ - control_dt: Optional[float] + control_dt: float action_space: Optional[gym.Space] _action: Optional[SpaceDictRecursive] @@ -27,7 +27,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: :param kwargs: Extra keyword arguments. See 'args'. """ # Define some attributes - self.control_dt = None + self.control_dt = 0.0 self.action_space = None self._action = None @@ -104,7 +104,7 @@ def compute_reward_terminal(self, info: Dict[str, Any]) -> float: class ObserveInterface: """Observer interface for both observers and environments. """ - observe_dt: Optional[float] + observe_dt: float observation_space: Optional[gym.Space] _observation: Optional[SpaceDictRecursive] @@ -118,14 +118,14 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: :param kwargs: Extra keyword arguments. See 'args'. """ # Define some attributes - self.observe_dt = None + self.observe_dt = 0.0 self.observation_space = None self._observation = None # Call super to allow mixing interfaces through multiple inheritance super().__init__(*args, **kwargs) # type: ignore[call-arg] - def fetch_observation(self) -> None: + def refresh_observation(self) -> None: """Refresh the observation. .. warning:: @@ -218,7 +218,7 @@ def _send_command(self, # breakpoint. A dedicated `BaseObserverBlock` must be used if one wants # to observe a low-frequency features instead of overloading # `compute_observation` and `_refresh_observation_space` directly. - self.fetch_observation() + self.refresh_observation() # Compute the command to send to the motors np.copyto(u_command, self.compute_command( diff --git a/python/gym_jiminy/common/gym_jiminy/common/pipeline_bases.py b/python/gym_jiminy/common/gym_jiminy/common/pipeline_bases.py index 473cdd07d..ae0ae18ab 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/pipeline_bases.py +++ b/python/gym_jiminy/common/gym_jiminy/common/pipeline_bases.py @@ -18,8 +18,10 @@ import numpy as np import gym +from jiminy_py.controller import ControllerHandleType + from .utils import ( - _is_breakpoint, _clamp, zeros, set_value, register_variables, + _is_breakpoint, _clamp, zeros, fill, set_value, register_variables, SpaceDictRecursive) from .generic_bases import ObserveAndControlInterface from .env_bases import BaseJiminyEnv @@ -54,14 +56,15 @@ def __init__(self, the wrapped block. Optional: disable by default. """ - # Initialize base wrapper and interfaces through multiple inheritance - super().__init__(env) - # Backup some user arguments self.augment_observation = augment_observation + # Initialize base wrapper and interfaces through multiple inheritance + super().__init__(env) + # Define some internal buffers self._dt_eps: Optional[float] = None + self._command = zeros(self.env.unwrapped.action_space) def __dir__(self) -> Sequence[str]: """Attribute lookup. @@ -101,28 +104,35 @@ def get_observation(self, bypass: bool = False) -> SpaceDictRecursive: return _clamp(self.observation_space, self._observation) def reset(self, - controller_hook: Optional[Callable[[], None]] = None, + controller_hook: Optional[ + Callable[[], Optional[ControllerHandleType]]] = None, **kwargs: Any) -> SpaceDictRecursive: """Reset the unified environment. In practice, it resets the environment and initializes the generic pipeline internal buffers through the use of 'controller_hook'. - :param controller_hook: Custom controller hook to use in place of the - one provided by the controller itself. Used - for chaining multiple `ControlledJiminyEnv`. - It is not meant to be defined manually. + :param controller_hook: Used internally for chaining multiple + `BasePipelineWrapper`. It is not meant to be + defined manually. Optional: None by default. :param kwargs: Extra keyword arguments to comply with OpenAI Gym API. """ # pylint: disable=unused-argument # Define chained controller hook - def register() -> None: + def register() -> ControllerHandleType: + """Register the block to the higher-level block. + + This method is used internally to make sure that `_setup` method + of connected blocks are called in the right order, namely from the + lowest-level block to the highest-level one, right after reset of + the low-level simulator and just before performing the first step. + """ nonlocal self, controller_hook # Assertion(s) for type checker - assert self.env is not None and self.env.simulator is not None + assert self.env is not None # Get the temporal resolution of simulator steps engine_options = self.simulator.engine.get_options() @@ -131,16 +141,17 @@ def register() -> None: # Initialize the pipeline wrapper self._setup() - # Register the controller handle or use the custom hook is defined - if controller_hook is None: - self.env.simulator.controller.set_controller_handle( - self._send_command) - else: - controller_hook() + # Forward the controller handle provided by the controller hook, + # if any, or use the one of the controller otherwise. + controller_handle = None + if controller_hook is not None: + controller_handle = controller_hook() + if controller_handle is None: + controller_handle = self._send_command + return controller_handle # Reset base pipeline - self.env.reset( # type: ignore[call-arg] - controller_hook=register, **kwargs) + self.env.reset(register, **kwargs) # type: ignore[call-arg] return self.get_observation() @@ -169,12 +180,41 @@ def step(self, def _setup(self) -> None: """Configure the wrapper. - This method does nothing by default. One is expected to overload it. + By default, it only resets some internal buffers. .. note:: This method must be called once, after the environment has been reset. This is done automatically when calling `reset` method. """ + fill(self._action, 0.0) + fill(self._command, 0.0) + fill(self._observation, 0.0) + + def compute_observation(self # type: ignore[override] + ) -> SpaceDictRecursive: + """Compute the unified observation. + + By default, it forwards the observation computed by the environment. + + :param measure: Observation of the environment. + """ + # pylint: disable=arguments-differ + + self.env.refresh_observation() + return self.get_observation(bypass=True).copy() # No deepcopy ! + + def compute_command(self, + measure: SpaceDictRecursive, + action: SpaceDictRecursive + ) -> SpaceDictRecursive: + """Compute the motors efforts to apply on the robot. + + By default, it forwards the command computed by the environment. + + :param measure: Observation of the environment. + :param action: Target to achieve. + """ + return self.env.compute_command(measure, action) class ControlledJiminyEnv(BasePipelineWrapper): @@ -230,7 +270,8 @@ class ControlledJiminyEnv(BasePipelineWrapper): def __init__(self, env: Union[gym.Wrapper, BaseJiminyEnv], controller: BaseControllerBlock, - augment_observation: bool = False): + augment_observation: bool = False, + **kwargs: Any): """ .. note:: As a reminder, `env.step_dt` refers to the learning step period, @@ -256,6 +297,8 @@ def __init__(self, if the observation space is of type `gym.spaces.Dict`. Optional: disable by default. + :param kwargs: Extra keyword arguments to allow automatic pipeline + wrapper generation. """ # Initialize base wrapper super().__init__(env, augment_observation) @@ -263,27 +306,14 @@ def __init__(self, # Backup user arguments self.controller = controller - # Define some internal buffers - self._target: Optional[SpaceDictRecursive] = None - self._command: Optional[np.ndarray] = None - self.controller_name: Optional[str] = None - - def _setup(self) -> None: - # Assertion(s) for type checker - assert self.simulator is not None - # Assertion(s) for type checker - assert (self.env.control_dt is not None and - self.env.action_space is not None and + assert (self.env.action_space is not None and self.env.observation_space is not None) # Reset the controller self.controller.reset(self.env) self.control_dt = self.controller.control_dt - # Assertion(s) for type checker - assert self.control_dt is not None - # Make sure the controller period is lower than environment timestep assert self.control_dt <= self.env.unwrapped.step_dt, ( "The control update period must be lower than or equal to the " @@ -298,16 +328,6 @@ def _setup(self) -> None: # Assertion(s) for type checker assert self.action_space is not None - # Initialize the controller's input action and output target - self._action = zeros(self.action_space) - self._target = zeros(self.env.action_space) - - # Initialize the unified observation with zero target - self._observation = self.compute_observation() - - # Initialize the command to apply on the robot - self._command = zeros(self.env.unwrapped.action_space) - # Check that 'augment_observation' can be enabled assert not self.augment_observation or isinstance( self.env.observation_space, gym.spaces.Dict), ( @@ -321,10 +341,25 @@ def _setup(self) -> None: 'targets', gym.spaces.Dict()).spaces[self.controller_name] = \ self.controller.action_space + # Initialize some internal buffers + self._action = zeros(self.action_space) + self._target = zeros(self.env.action_space) + self._observation = zeros(self.observation_space) + + def _setup(self) -> None: + """Configure the wrapper. + + In addition to the base implementation, it resgisters the controller's + target to the telemetry. + """ + # Call base implementation + super()._setup() + + # Reset some additional internal buffers + fill(self._target, 0.0) + # Register the controller target to the telemetry. - # It may be useful later for computing the terminal reward or debug. - # Note that it is not necessary for the controller to be fully - # initialized before registering variables. + # It may be useful for computing the terminal reward or debugging. register_variables( self.simulator.controller, self.controller.get_fieldnames(), self._action, self.controller_name) @@ -344,14 +379,12 @@ def compute_command(self, :param measure: Observation of the environment. :param action: High-level target to achieve. """ - # Assertion(s) for type checker - assert self.simulator is not None and self._command is not None - # Update the target to send to the subsequent block if necessary. # Note that `_observation` buffer has already been updated right before # calling this method by `_send_command`, so it can be used as measure # argument without issue. - if _is_breakpoint(measure['t'], self.control_dt, self._dt_eps): + t = self.simulator.stepper_state.t + if _is_breakpoint(t, self.control_dt, self._dt_eps): target = self.controller.compute_command(self._observation, action) set_value(self._target, target) @@ -360,8 +393,11 @@ def compute_command(self, # update the command of the right period. Ultimately, this is done # automatically by the engine, which is calling `_send_command` at the # right period. - np.copyto(self._command, self.env.compute_command( - self._observation, self._target)) + if self.env.simulator.is_simulation_running: + # Do not update command during the first iteration because the + # action is undefined at this point + np.copyto(self._command, self.env.compute_command( + self._observation, self._target)) return self._command @@ -382,13 +418,14 @@ def compute_observation(self # type: ignore[override] :returns: Original environment observation, eventually including controllers targets if requested. """ - # pylint: disable=arguments-differ + # Get environment observation + obs = super().compute_observation() - self.env.fetch_observation() - obs = self.get_observation(bypass=True).copy() # No deepcopy ! + # Add target to observation if requested if self.augment_observation: obs.setdefault('targets', OrderedDict())[ self.controller_name] = self._action + return obs def step(self, @@ -439,7 +476,8 @@ class ObservedJiminyEnv(BasePipelineWrapper): def __init__(self, env: Union[gym.Wrapper, BaseJiminyEnv], observer: BaseObserverBlock, - augment_observation: bool = False): + augment_observation: bool = False, + **kwargs: Any): """ :param env: Environment to control. It can be an already controlled environment wrapped in `ObservedJiminyEnv` if one desires @@ -451,6 +489,8 @@ def __init__(self, option is only available if the observation space is of type `gym.spaces.Dict`. Optional: disable by default. + :param kwargs: Extra keyword arguments to allow automatic pipeline + wrapper generation. """ # Initialize base wrapper super().__init__(env, augment_observation) @@ -458,63 +498,6 @@ def __init__(self, # Backup user arguments self.observer = observer - # Reset the unified environment - self.reset() - - def compute_observation(self # type: ignore[override] - ) -> SpaceDictRecursive: - """Compute high-level features based on the current wrapped - environment's observation. - - It gathers the original observation from the environment with the - features computed by the observer, if requested, otherwise it forwards - the features directly without any further processing. - - .. warning:: - Beware it updates and returns '_observation' buffer to deal with - multiple observers with different update periods. Even so, it is - safe to call this method multiple times successively. - - :returns: Updated part of the observation only for efficiency. - """ - # pylint: disable=arguments-differ - - # Refresh environment observation - self.env.fetch_observation() - if self.augment_observation: - obs = self.get_observation(bypass=True).copy() # No deepcopy ! - - # Get the current time - t = self.simulator.stepper_state.t - - # Update observed features if necessary - if _is_breakpoint(t, self.observe_dt, self._dt_eps): - features = self.observer.compute_observation(obs) - if self.augment_observation: - obs.setdefault( - 'features', OrderedDict())[self.observer_name] = features - else: - obs = features - else: - if not self.augment_observation: - obs = OrderedDict() # Nothing new to observe. - - return obs - - def compute_command(self, - measure: SpaceDictRecursive, - action: SpaceDictRecursive - ) -> SpaceDictRecursive: - """Compute the motors efforts to apply on the robot. - - In practice, it forwards the command computed by the environment. - - :param measure: Observation of the environment. - :param action: Target to achieve. - """ - return self.env.compute_command(measure, action) - - def _setup(self) -> None: # Assertion(s) for type checker assert (self.env.action_space is not None and self.env.observation_space is not None) @@ -525,13 +508,6 @@ def _setup(self) -> None: # Update the action space self.action_space = self.env.action_space - # Initialize the unified observation with zero target - self._observation = self.compute_observation() - - # Initialize the environment's action and command - self._action = zeros(self.action_space) - self._command = zeros(self.env.unwrapped.action_space) - # Reset the observer self.observer.reset(self.env) @@ -557,6 +533,46 @@ def _setup(self) -> None: else: self.observation_space = self.observer.observation_space + # Initialize some internal buffers + self._action = zeros(self.action_space) + self._observation = zeros(self.observation_space) + + def compute_observation(self # type: ignore[override] + ) -> SpaceDictRecursive: + """Compute high-level features based on the current wrapped + environment's observation. + + It gathers the original observation from the environment with the + features computed by the observer, if requested, otherwise it forwards + the features directly without any further processing. + + .. warning:: + Beware it updates and returns '_observation' buffer to deal with + multiple observers with different update periods. Even so, it is + safe to call this method multiple times successively. + + :returns: Updated part of the observation only for efficiency. + """ + # pylint: disable=arguments-differ + + # Get environment observation + obs = super().compute_observation() + + # Update observed features if necessary + t = self.simulator.stepper_state.t + if _is_breakpoint(t, self.observe_dt, self._dt_eps): + features = self.observer.compute_observation(obs) + if self.augment_observation: + obs.setdefault( + 'features', OrderedDict())[self.observer_name] = features + else: + obs = features + else: + if not self.augment_observation: + obs = OrderedDict() # Nothing new to observe. + + return obs + def build_pipeline(env_config: Tuple[ Type[BaseJiminyEnv], diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils.py b/python/gym_jiminy/common/gym_jiminy/common/utils.py index 49d5e4cb4..61a4535af 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils.py @@ -33,23 +33,37 @@ def zeros(space: gym.Space) -> SpaceDictRecursive: f"Space of type {type(space)} is not supported by this method.") +def fill(data: SpaceDictRecursive, + fill_value: float) -> None: + """Set every element of 'data' from `Gym.Space` to scalar 'fill_value'. + """ + if isinstance(data, dict): + for sub_data in data.values(): + fill(sub_data, fill_value) + elif isinstance(data, np.ndarray): + data.fill(fill_value) + else: + raise NotImplementedError( + f"Data of type {type(data)} is not supported by this method.") + + def set_value(data: SpaceDictRecursive, - fill_value: SpaceDictRecursive) -> None: - """Partially set 'data' from `Gym.Space` to 'fill_value'. + value: SpaceDictRecursive) -> None: + """Partially set 'data' from `Gym.Space` to 'value'. It avoids memory allocation, so that memory pointers of 'data' remains unchanged. A direct consequences, it is necessary to preallocate memory beforehand, and to work with fixed size buffers. .. note:: - If 'data' is a dictionary, 'fill_value' must be a subtree of 'data', + If 'data' is a dictionary, 'value' must be a subtree of 'data', whose leaf values must be broadcastable with the ones of 'data'. """ if isinstance(data, dict): - for field, sub_val in fill_value.items(): + for field, sub_val in value.items(): set_value(data[field], sub_val) elif isinstance(data, np.ndarray): - np.copyto(data, fill_value) + np.copyto(data, value) else: raise NotImplementedError( f"Data of type {type(data)} is not supported by this method.") diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers.py index bf0f5b88d..fce8ca436 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers.py @@ -1,13 +1,13 @@ from copy import deepcopy from functools import reduce from collections import deque -from typing import Tuple, Type, Dict, Sequence, List, Optional, Any, Iterator +from typing import Tuple, Type, Dict, Sequence, List, Any, Iterator import numpy as np import gym -from .utils import SpaceDictRecursive +from .utils import zeros, SpaceDictRecursive class PartialFrameStack(gym.Wrapper): @@ -30,46 +30,7 @@ def __init__(self, will be stacked. :param num_stack: Number of observation frames to partially stack. """ - # Backup user arguments - self.nested_fields_list: List[List[str]] = list( - map(list, nested_fields_list)) # type: ignore[arg-type] - self.leaf_fields_list: List[List[str]] = [] - self.num_stack = num_stack - - # Define some internal buffers - self._observation: Optional[SpaceDictRecursive] = None - - # Initialize base wrapper - super().__init__(env) - - # Create internal buffers - self._frames: List[deque] = [] - - def get_observation(self) -> SpaceDictRecursive: - assert (self._observation is not None and - all(len(frames) == self.num_stack for frames in self._frames)) - - # Replace nested fields of original observation by the stacked ones - for fields, frames in zip(self.leaf_fields_list, self._frames): - root_obs = reduce(lambda d, key: d[key], fields[:-1], - self._observation) - root_obs[fields[-1]] = np.stack(frames) - - return self._observation - - def step(self, - action: SpaceDictRecursive - ) -> Tuple[SpaceDictRecursive, float, bool, Dict[str, Any]]: - self._observation, reward, done, info = self.env.step(action) - - # Backup the nested observation fields to stack - for fields, frames in zip(self.leaf_fields_list, self._frames): - leaf_obs = reduce(lambda d, key: d[key], fields, self._observation) - frames.append(leaf_obs) - - return self.get_observation(), reward, done, info - - def reset(self, **kwargs: Any) -> SpaceDictRecursive: + # Define helper that will be used to determine the leaf fields to stack def _get_branches(root: Any) -> Iterator[List[str]]: if isinstance(root, dict): for field, node in root.items(): @@ -79,10 +40,20 @@ def _get_branches(root: Any) -> Iterator[List[str]]: else: yield [field] - self._observation = self.env.reset(**kwargs) + # Backup user arguments + self.nested_fields_list: List[List[str]] = list( + map(list, nested_fields_list)) # type: ignore[arg-type] + self.num_stack = num_stack + + # Initialize base wrapper + super().__init__(env) + + # Define some internal buffers + self._observation: SpaceDictRecursive = zeros( + self.env.observation_space) - # Determine leaf fields to stack - self.leaf_fields_list = [] + # Get the leaf fields to stack + self.leaf_fields_list: List[List[str]] = [] for fields in self.nested_fields_list: root_field = reduce( lambda d, key: d[key], fields, self._observation) @@ -109,9 +80,34 @@ def _get_branches(root: Any) -> Iterator[List[str]]: root_space.spaces[fields[-1]] = gym.spaces.Box( low=low, high=high, dtype=space.dtype) - # Allocate the frames buffers - self._frames = [deque(maxlen=self.num_stack) - for _ in range(len(self.leaf_fields_list))] + # Allocate internal frames buffers + self._frames: List[deque] = [ + deque(maxlen=self.num_stack) + for _ in range(len(self.leaf_fields_list))] + + def get_observation(self) -> SpaceDictRecursive: + # Replace nested fields of original observation by the stacked ones + for fields, frames in zip(self.leaf_fields_list, self._frames): + root_obs = reduce(lambda d, key: d[key], fields[:-1], + self._observation) + root_obs[fields[-1]] = np.stack(frames) + return self._observation + + def step(self, + action: SpaceDictRecursive + ) -> Tuple[SpaceDictRecursive, float, bool, Dict[str, Any]]: + # Perform a single step + self._observation, reward, done, info = self.env.step(action) + + # Backup the nested observation fields to stack + for fields, frames in zip(self.leaf_fields_list, self._frames): + leaf_obs = reduce(lambda d, key: d[key], fields, self._observation) + frames.append(leaf_obs) + + return self.get_observation(), reward, done, info + + def reset(self, **kwargs: Any) -> SpaceDictRecursive: + self._observation = self.env.reset(**kwargs) # Initialize the frames by duplicating the original one for fields, frames in zip(self.leaf_fields_list, self._frames): @@ -161,11 +157,21 @@ def build_wrapper(env_config: Tuple[ (wrapper_class,), {}) - def __init__(self: wrapped_env_class) -> None: # type: ignore[valid-type] + def __init__(self: gym.Wrapper) -> None: env = env_class(**env_kwargs) super(wrapped_env_class, self).__init__( # type: ignore[arg-type] env, **wrapper_kwargs) + def __dir__(self: gym.Wrapper) -> Sequence[str]: + """Attribute lookup. + + It is mainly used by autocomplete feature of Ipython. It is overloaded + to get consistent autocompletion wrt `getattr`. + """ + return super( # type: ignore[arg-type] + wrapped_env_class, self).__dir__() + self.env.__dir__() + wrapped_env_class.__init__ = __init__ # type: ignore[misc] + wrapped_env_class.__dir__ = __dir__ # type: ignore[assignment] return wrapped_env_class diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py b/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py index ac7c30f6d..5937b7c0b 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py @@ -63,11 +63,11 @@ def __init__(self, debug: bool = False, **kwargs): debug=debug, **kwargs) - def _refresh_observation_space(self) -> None: - self.observation_space = self._get_state_space() + # def _refresh_observation_space(self) -> None: + # self.observation_space = self._get_state_space() - def fetch_obs(self) -> None: - return np.concatenate(self._state) + # def compute_observation(self) -> None: + # return np.concatenate(self._state) ANYmalPDControlJiminyEnv = build_pipeline(**{ diff --git a/python/gym_jiminy/unit_py/atlas_standing_meshcat.png b/python/gym_jiminy/unit_py/atlas_standing_meshcat.png index fca496b08..28f252a79 100644 Binary files a/python/gym_jiminy/unit_py/atlas_standing_meshcat.png and b/python/gym_jiminy/unit_py/atlas_standing_meshcat.png differ diff --git a/python/gym_jiminy/unit_py/test_pipeline_control.py b/python/gym_jiminy/unit_py/test_pipeline_control.py index f787b9f4b..c5dc80a23 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_control.py +++ b/python/gym_jiminy/unit_py/test_pipeline_control.py @@ -16,13 +16,22 @@ def setUp(self): """ TODO: Write documentation """ warnings.filterwarnings("ignore", category=DeprecationWarning) - self.env = AtlasPDControlJiminyEnv() + self.env = AtlasPDControlJiminyEnv(debug=False) def test_pid_standing(self): """ TODO: Write documentation """ + # Check that it is not possible to get simulation log at this point + self.assertRaises(RuntimeError, self.env.get_log) + # Reset the environment - obs_init = self.env.reset() + def configure_telemetry() -> None: + nonlocal self + engine_options = self.env.simulator.engine.get_options() + engine_options['telemetry']['enableVelocity'] = True + self.env.simulator.engine.set_options(engine_options) + + obs_init = self.env.reset(controller_hook=configure_telemetry) # The initial target corresponds to the initial joints state, so that # the robot stand-still. @@ -59,8 +68,8 @@ def test_pid_standing(self): data = log_data[log_name] self.assertTrue(np.all(np.abs(data) < 1e-9)) - # Check that the velocity is close to zero at the end + # Check that the whole-body robot velocity is close to zero at the end velocity_mes = np.stack([ - log_data['.'.join((encoder.type, name, encoder.fieldnames[1]))] - for name in self.env.robot.sensors_names[encoder.type]], axis=-1) + log_data['.'.join(('HighLevelController', name))] + for name in self.env.robot.logfile_velocity_headers], axis=-1) self.assertTrue(np.all(np.abs(velocity_mes[-1000:]) < 1e-3)) diff --git a/python/jiminy_py/src/jiminy_py/controller.py b/python/jiminy_py/src/jiminy_py/controller.py index 75d2b3447..c26e9bfb9 100644 --- a/python/jiminy_py/src/jiminy_py/controller.py +++ b/python/jiminy_py/src/jiminy_py/controller.py @@ -12,6 +12,11 @@ from tqdm import tqdm +ControllerHandleType = Callable[[ + float, np.ndarray, np.ndarray, jiminy.sensorsData, np.ndarray +], None] + + class BaseJiminyController(jiminy.ControllerFunctor): """Base class to instantiate a Jiminy controller based on a callable function that can be changed on-the-fly. @@ -63,12 +68,9 @@ def reset(self) -> None: super().reset() self.close_progress_bar() - def set_controller_handle( - self, - controller_handle: Callable[[ - float, np.ndarray, np.ndarray, jiminy.sensorsData, np.ndarray - ], None] - ) -> None: + def set_controller_handle(self, + controller_handle: ControllerHandleType + ) -> None: r"""Set the controller callback function to use. :param compute_command: diff --git a/python/jiminy_py/src/jiminy_py/dynamics.py b/python/jiminy_py/src/jiminy_py/dynamics.py index 87ea317ed..484e31e2d 100644 --- a/python/jiminy_py/src/jiminy_py/dynamics.py +++ b/python/jiminy_py/src/jiminy_py/dynamics.py @@ -496,7 +496,7 @@ def compute_freeflyer_state_from_fixed_body( w_M_ff = w_M_ground.act(ff_M_fixed_body.inverse()) position[:7] = pin.SE3ToXYZQUAT(w_M_ff) - if fixed_body_name is None: + if fixed_body_name is not None: if velocity is not None: ff_v_fixed_body = get_body_world_velocity( robot, fixed_body_name, use_theoretical_model=False)