diff --git a/CMakeLists.txt b/CMakeLists.txt index 13800f9ea..16d4ab678 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.12.4) # Set the build version (specify a tweak version to indicated post-release if needed) -set(BUILD_VERSION 1.8.5.1) +set(BUILD_VERSION 1.8.6) # MSVC runtime library flags are defined by 'CMAKE_MSVC_RUNTIME_LIBRARY' if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.15.7) diff --git a/INSTALL.md b/INSTALL.md index 9bb43f892..192129f2d 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -116,7 +116,7 @@ make install -j2 ### Prerequisites -You have to preinstall by yourself the (free) MSVC 2019 toolchain. +You have to preinstall by yourself the (free) MSVC 2019+ toolchain. Then, install `setuptools`, `wheel` and `numpy`. @@ -152,24 +152,23 @@ if (-not (Test-Path -PathType Container $RootDir/build)) { New-Item -ItemType "directory" -Force -Path "$RootDir/build" } Set-Location -Path $RootDir/build -cmake "$RootDir" -G "Visual Studio 16 2019" -T "v142" -DCMAKE_GENERATOR_PLATFORM=x64 ` - -DCMAKE_INSTALL_PREFIX="$InstallDir" -DCMAKE_MODULE_PATH="$InstallDir" ` +cmake "$RootDir" -G "Visual Studio 17 2022" -DCMAKE_GENERATOR_PLATFORM=x64 ` + -DCMAKE_MSVC_RUNTIME_LIBRARY="MultiThreaded$<$:Debug>DLL" ` + -DCMAKE_INSTALL_PREFIX="$InstallDir" -DCMAKE_PREFIX_PATH="$InstallDir" ` + -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF -DCMAKE_VERBOSE_MAKEFILE=ON ` -DBOOST_ROOT="$InstallDir" -DBoost_INCLUDE_DIR="$InstallDir/include" ` - -DBoost_NO_SYSTEM_PATHS=TRUE -DBoost_NO_BOOST_CMAKE=TRUE ` - -DBoost_USE_STATIC_LIBS=OFF ` - -DBUILD_TESTING=ON -DBUILD_EXAMPLES=ON -DBUILD_PYTHON_INTERFACE=ON ` - -DCMAKE_CXX_FLAGS="-DBOOST_ALL_NO_LIB -DBOOST_LIB_DIAGNOSTIC -DBOOST_CORE_USE_GENERIC_CMATH" -cmake --build . --target all --config "${env:BUILD_TYPE}" --parallel 8 - -if (-not (Test-Path -PathType Container "$RootDir/build/PyPi/jiminy_py/src/jiminy_py/core")) { - New-Item -ItemType "directory" -Force -Path "$RootDir/build/PyPi/jiminy_py/src/jiminy_py/core" + -DBoost_NO_SYSTEM_PATHS=TRUE -DBoost_NO_BOOST_CMAKE=TRUE -DBoost_USE_STATIC_LIBS=ON ` + -DBUILD_TESTING=ON -DBUILD_EXAMPLES=ON -DBUILD_PYTHON_INTERFACE=ON `` + -DCMAKE_CXX_FLAGS="${env:CMAKE_CXX_FLAGS} $( + ) -DBOOST_ALL_NO_LIB -DBOOST_LIB_DIAGNOSTIC -DBOOST_CORE_USE_GENERIC_CMATH $( + ) -DEIGENPY_STATIC -DURDFDOM_STATIC -DHPP_FCL_STATIC -DPINOCCHIO_STATIC" +cmake --build . --target ALL_BUILD --config "${env:BUILD_TYPE}" --parallel 2 + +if (-not (Test-Path -PathType Container "$RootDir/build/pypi/jiminy_py/src/jiminy_py")) { + New-Item -ItemType "directory" -Force -Path "$RootDir/build/pypi/jiminy_py/src/jiminy_py/core" } -Copy-Item -Path "$InstallDir/lib/boost_numpy*.dll" ` - -Destination "$RootDir/build/PyPi/jiminy_py/src/jiminy_py/core" -Copy-Item -Path "$InstallDir/lib/boost_python*.dll" ` - -Destination "$RootDir/build/PyPi/jiminy_py/src/jiminy_py/core" -Copy-Item -Path "$InstallDir/lib/site-packages/*" ` - -Destination "$RootDir/build/PyPi/jiminy_py/src/jiminy_py" -Recurse - -cmake --build . --target install --config "${env:BUILD_TYPE}" +Copy-Item -Force -Recurse -Path "$InstallDir/lib/site-packages/*" ` + -Destination "$RootDir/build/pypi/jiminy_py/src/jiminy_py/core" + +cmake --build . --target INSTALL --config "${env:BUILD_TYPE}" ``` diff --git a/README.md b/README.md index 7fab1181a..593f2d418 100644 --- a/README.md +++ b/README.md @@ -29,11 +29,12 @@ Beside a strong focus on performance to answer machine learning's need for runni ### General -- Simulation of multi-body systems using minimal coordinates and Lagrangian dynamics. +- Simulation of multi-body systems using **minimal coordinates and Lagrangian dynamics**. - Comprehensive API for computing dynamic quantities and their derivatives thanks to Pinocchio. - C++ core with full python bindings, providing user API parity between both languages. - Designed with machine learning in mind, with seamless wrapping of robots as [OpenAI Gym](https://github.com/openai/gym) environments using one-liners. Jiminy provides both the physical engine and the robot model (including sensors) required for learning. -- Rich simulation log output, easily customizable for recording, introspection and debugging. The simulation log is made available in RAM directly for fast access, and can be exported in raw binary or [HDF5](./docs/spec/src/tlmc_format_specification.md) format. +- **Fully modular learning environment pipeline design**, with plug-and-play robotic observers and controllers, but also composable reward components and termination conditions. See example [here](./python/gym_jiminy/unit_py/data/anymal_pipeline.toml). +- Rich simulation log output, easily customizable for recording, introspection and debugging. The simulation log is made available in RAM directly for fast access, and can be exported in binary or [HDF5](./docs/spec/src/tlmc_format_specification.md) format. - Dedicated integration in Google Colab, Jupyter Lab, Mybinder and VSCode working out-of-the-box - including interactive 3D viewer based on [Meshcat](https://github.com/rdeits/MeshCat.jl). This facilitates working on remote environments. - Synchronous and single-threaded offscreen rendering capability, GPU-accelerated without graphical server, based on [Panda3d](https://github.com/panda3d/panda3d). - Easy to install: `pip` is all that is needed to [get you started](#getting-started) ! Support Linux, Mac and Windows platforms. @@ -43,7 +44,7 @@ Beside a strong focus on performance to answer machine learning's need for runni - Provide both classical phenomenological force-level spring-damper contact model and constraint solver satisfying the maximum energy dissipation principle. - Support contact and collision with the ground from a fixed set of contact points (primitives and meshes yet to come). - Simulate multiple articulated systems simultaneously interacting with each other, to support use cases such as multi-agent learning or swarm robotics. -- Compliant spherical joints with spring-damper dynamics to model mechanical deformation, a common phenomenon particularly in legged robotics. +- Modeling of motor transmission backlash and structural deformation via compliant spherical joints with spring-damper dynamics, which are common phenomena, particularly in bipedal robotics. - Simulate both continuous or discrete-time controller, with possibly different controller and sensor update frequencies. A more complete list of features is available on the [wiki](https://github.com/duburcqa/jiminy/wiki). diff --git a/core/examples/double_pendulum/double_pendulum.cc b/core/examples/double_pendulum/double_pendulum.cc index ea11d4c04..5b6cf434c 100644 --- a/core/examples/double_pendulum/double_pendulum.cc +++ b/core/examples/double_pendulum/double_pendulum.cc @@ -66,12 +66,10 @@ int main(int /* argc */, char * /* argv */[]) // Instantiate and configuration the robot std::vector motorJointNames{"SecondPendulumJoint"}; - auto robot = std::make_shared(); GenericConfig modelOptions = robot->getModelOptions(); GenericConfig & jointsOptions = boost::get(modelOptions.at("joints")); boost::get(jointsOptions.at("positionLimitFromUrdf")) = true; - boost::get(jointsOptions.at("velocityLimitFromUrdf")) = true; robot->setModelOptions(modelOptions); robot->initialize(urdfPath.string(), false, {dataPath.string()}); for (const std::string & jointName : motorJointNames) @@ -79,6 +77,9 @@ int main(int /* argc */, char * /* argv */[]) auto motor = std::make_shared(jointName); robot->attachMotor(motor); motor->initialize(jointName); + GenericConfig motorOptions = motor->getOptions(); + boost::get(motorOptions.at("velocityLimitFromUrdf")) = true; + motor->setOptions(motorOptions); } // Instantiate and configuration the controller @@ -95,7 +96,7 @@ int main(int /* argc */, char * /* argv */[]) boost::get(telemetryOptions.at("enableAcceleration")) = true; boost::get(telemetryOptions.at("enableForceExternal")) = false; boost::get(telemetryOptions.at("enableCommand")) = true; - boost::get(telemetryOptions.at("enableMotorEffort")) = true; + boost::get(telemetryOptions.at("enableEffort")) = true; boost::get(telemetryOptions.at("enableEnergy")) = true; GenericConfig & worldOptions = boost::get(simuOptions.at("world")); boost::get(worldOptions.at("gravity"))[2] = -9.81; diff --git a/core/examples/external_project/double_pendulum.cc b/core/examples/external_project/double_pendulum.cc index fa40433db..8809fb521 100644 --- a/core/examples/external_project/double_pendulum.cc +++ b/core/examples/external_project/double_pendulum.cc @@ -67,7 +67,6 @@ int main(int argc, char * argv[]) GenericConfig modelOptions = robot->getModelOptions(); GenericConfig & jointsOptions = boost::get(modelOptions.at("joints")); boost::get(jointsOptions.at("positionLimitFromUrdf")) = true; - boost::get(jointsOptions.at("velocityLimitFromUrdf")) = true; robot->setModelOptions(modelOptions); robot->initialize(urdfPath.string(), false, {}); @@ -75,9 +74,12 @@ int main(int argc, char * argv[]) auto motor = std::make_shared("motor"); robot->attachMotor(motor); motor->initialize("SecondPendulumJoint"); + GenericConfig motorOptions = motor->getOptions(); + boost::get(motorOptions.at("velocityLimitFromUrdf")) = true; + motor->setOptions(motorOptions); auto sensor = std::make_shared("encoder"); robot->attachSensor(sensor); - sensor->initialize("SecondPendulumJoint"); + sensor->initialize("SecondPendulumJoint", true); // Print encoder sensor index std::cout << "Encoder sensor index: " << sensor->getIndex() << std::endl; diff --git a/core/include/jiminy/core/engine/engine.h b/core/include/jiminy/core/engine/engine.h index 7b202a534..ea3716658 100644 --- a/core/include/jiminy/core/engine/engine.h +++ b/core/include/jiminy/core/engine/engine.h @@ -146,6 +146,7 @@ namespace jiminy Eigen::VectorXd command{}; Eigen::VectorXd u{}; Eigen::VectorXd uMotor{}; + Eigen::VectorXd uTransmission{}; Eigen::VectorXd uInternal{}; Eigen::VectorXd uCustom{}; ForceVector fExternal{}; @@ -198,9 +199,9 @@ namespace jiminy std::vector logPositionFieldnames{}; std::vector logVelocityFieldnames{}; std::vector logAccelerationFieldnames{}; + std::vector logEffortFieldnames{}; std::vector logForceExternalFieldnames{}; std::vector logCommandFieldnames{}; - std::vector logMotorEffortFieldnames{}; std::string logEnergyFieldname{}; /// \brief Internal buffer with the state for the integration loop. @@ -283,15 +284,6 @@ namespace jiminy return config; }; - GenericConfig getDefaultJointOptions() - { - GenericConfig config; - config["boundStiffness"] = 1.0e7; - config["boundDamping"] = 1.0e4; - - return config; - }; - GenericConfig getDefaultWorldOptions() { GenericConfig config; @@ -337,9 +329,9 @@ namespace jiminy config["enableVelocity"] = true; config["enableAcceleration"] = true; config["enableForceExternal"] = false; - config["enableCommand"] = true; - config["enableMotorEffort"] = true; - config["enableEnergy"] = true; + config["enableCommand"] = false; + config["enableEffort"] = false; + config["enableEnergy"] = false; return config; }; @@ -349,7 +341,6 @@ namespace jiminy config["telemetry"] = getDefaultTelemetryOptions(); config["stepper"] = getDefaultStepperOptions(); config["world"] = getDefaultWorldOptions(); - config["joints"] = getDefaultJointOptions(); config["constraints"] = getDefaultConstraintOptions(); config["contacts"] = getDefaultContactOptions(); @@ -394,18 +385,6 @@ namespace jiminy } }; - struct JointOptions - { - const double boundStiffness; - const double boundDamping; - - JointOptions(const GenericConfig & options) : - boundStiffness{boost::get(options.at("boundStiffness"))}, - boundDamping{boost::get(options.at("boundDamping"))} - { - } - }; - struct WorldOptions { const Eigen::VectorXd gravity; @@ -460,7 +439,7 @@ namespace jiminy const bool enableAcceleration; const bool enableForceExternal; const bool enableCommand; - const bool enableMotorEffort; + const bool enableEffort; const bool enableEnergy; TelemetryOptions(const GenericConfig & options) : @@ -470,7 +449,7 @@ namespace jiminy enableAcceleration{boost::get(options.at("enableAcceleration"))}, enableForceExternal{boost::get(options.at("enableForceExternal"))}, enableCommand{boost::get(options.at("enableCommand"))}, - enableMotorEffort{boost::get(options.at("enableMotorEffort"))}, + enableEffort{boost::get(options.at("enableEffort"))}, enableEnergy{boost::get(options.at("enableEnergy"))} { } @@ -481,7 +460,6 @@ namespace jiminy const TelemetryOptions telemetry; const StepperOptions stepper; const WorldOptions world; - const JointOptions joints; const ConstraintOptions constraints; const ContactOptions contacts; @@ -489,7 +467,6 @@ namespace jiminy telemetry{boost::get(options.at("telemetry"))}, stepper{boost::get(options.at("stepper"))}, world{boost::get(options.at("world"))}, - joints{boost::get(options.at("joints"))}, constraints{boost::get(options.at("constraints"))}, contacts{boost::get(options.at("contacts"))} { diff --git a/core/include/jiminy/core/hardware/abstract_motor.h b/core/include/jiminy/core/hardware/abstract_motor.h index f5437a8aa..329810793 100644 --- a/core/include/jiminy/core/hardware/abstract_motor.h +++ b/core/include/jiminy/core/hardware/abstract_motor.h @@ -24,8 +24,8 @@ namespace jiminy /// state of every motor. struct MotorSharedStorage { - /// \brief Buffer storing the current true motor efforts. - Eigen::VectorXd data_; + /// \brief Buffer storing the current motor efforts on joint and motor sides respectively. + Eigen::MatrixX2d data_; /// \brief Vector of pointers to the motors. std::vector motors_; /// \brief Number of motors @@ -43,9 +43,10 @@ namespace jiminy { GenericConfig config; config["mechanicalReduction"] = 1.0; - config["enableCommandLimit"] = true; - config["commandLimitFromUrdf"] = true; - config["commandLimit"] = 0.0; + config["velocityLimitFromUrdf"] = true; + config["velocityLimit"] = 0.0; + config["effortLimitFromUrdf"] = true; + config["effortLimit"] = 0.0; config["enableArmature"] = false; config["armature"] = 0.0; config["enableBacklash"] = false; @@ -58,9 +59,10 @@ namespace jiminy { /// \brief Mechanical reduction ratio of transmission (joint/motor), usually >= 1.0. const double mechanicalReduction; - const bool enableCommandLimit; - const bool commandLimitFromUrdf; - const double commandLimit; + const bool velocityLimitFromUrdf; + const double velocityLimit; + const bool effortLimitFromUrdf; + const double effortLimit; const bool enableArmature; const double armature; const bool enableBacklash; @@ -68,9 +70,10 @@ namespace jiminy AbstractMotorOptions(const GenericConfig & options) : mechanicalReduction(boost::get(options.at("mechanicalReduction"))), - enableCommandLimit(boost::get(options.at("enableCommandLimit"))), - commandLimitFromUrdf(boost::get(options.at("commandLimitFromUrdf"))), - commandLimit(boost::get(options.at("commandLimit"))), + velocityLimitFromUrdf(boost::get(options.at("velocityLimitFromUrdf"))), + velocityLimit(boost::get(options.at("velocityLimit"))), + effortLimitFromUrdf(boost::get(options.at("effortLimitFromUrdf"))), + effortLimit(boost::get(options.at("effortLimit"))), enableArmature(boost::get(options.at("enableArmature"))), armature(boost::get(options.at("armature"))), enableBacklash(boost::get(options.at("enableBacklash"))), @@ -84,7 +87,7 @@ namespace jiminy public: /// \param[in] name Name of the motor. - explicit AbstractMotorBase(const std::string & name) noexcept; + explicit AbstractMotorBase(const std::string & name); virtual ~AbstractMotorBase(); /// \brief Refresh the proxies. @@ -105,10 +108,11 @@ namespace jiminy const GenericConfig & getOptions() const noexcept; /// \brief Actual effort of the motor at the current time. - double get() const; + std::tuple get() const; /// \brief Actual effort of all the motors at the current time. - const Eigen::VectorXd & getAll() const; + std::tuple, Eigen::Ref> + getAll() const; /// \brief Set the configuration options of the motor. /// @@ -138,22 +142,22 @@ namespace jiminy /// \brief Index of the joint associated with the motor in the kinematic tree. pinocchio::JointIndex getJointIndex() const; - /// \brief Type of joint associated with the motor. - JointModelType getJointType() const; + /// \brief Maximum position of the actuated joint translated on motor side. + double getPositionLimitLower() const; - /// \brief Index of the joint associated with the motor in configuration vector. - Eigen::Index getJointPositionIndex() const; + /// \brief Minimum position of the actuated joint translated on motor side. + double getPositionLimitUpper() const; - /// \brief Index of the joint associated with the motor in the velocity vector. - Eigen::Index getJointVelocityIndex() const; + /// \brief Maximum velocity of the motor. + double getVelocityLimit() const; /// \brief Maximum effort of the motor. - double getCommandLimit() const; + double getEffortLimit() const; - /// \brief Rotor inertia of the motor. + /// \brief Rotor inertia of the motor on joint side. double getArmature() const; - /// \brief Backlash of the transmission. + /// \brief Backlash of the transmission on joint side. double getBacklash() const; /// \brief Request the motor to update its actual effort based of the input data. @@ -193,7 +197,7 @@ namespace jiminy protected: /// \brief Reference to the last data buffer corresponding to the true effort of the motor. - double & data(); + std::tuple data(); private: /// \brief Attach the sensor to a robot @@ -223,9 +227,10 @@ namespace jiminy std::string jointName_{}; pinocchio::JointIndex jointIndex_{0}; JointModelType jointType_{JointModelType::UNSUPPORTED}; - Eigen::Index jointPositionIndex_{0}; - Eigen::Index jointVelocityIndex_{0}; - double commandLimit_{0.0}; + double positionLimitLower_{}; + double positionLimitUpper_{}; + double velocityLimit_{0.0}; + double effortLimit_{0.0}; double armature_{0.0}; double backlash_{0.0}; diff --git a/core/include/jiminy/core/hardware/abstract_sensor.h b/core/include/jiminy/core/hardware/abstract_sensor.h index ad49b42fa..2300a752b 100644 --- a/core/include/jiminy/core/hardware/abstract_sensor.h +++ b/core/include/jiminy/core/hardware/abstract_sensor.h @@ -37,7 +37,7 @@ namespace jiminy boost::circular_buffer data_; /// \brief Current sensor measurements. Eigen::MatrixXd measurements_; - /// \brief Vector of pointers to the sensors. + /// \brief Vector of raw pointers to the sensors. std::vector sensors_; /// \brief Number of sensors currently sharing this buffer. std::size_t num_; @@ -106,7 +106,7 @@ namespace jiminy public: /// \param[in] name Name of the sensor - explicit AbstractSensorBase(const std::string & name) noexcept; + explicit AbstractSensorBase(const std::string & name); virtual ~AbstractSensorBase(); /// \brief Reset the internal state of the sensors. @@ -266,9 +266,9 @@ namespace jiminy /// to be filled up at this stage. /// /// \return Eigen Reference to a Eigen Vector corresponding to the last data recorded. - virtual Eigen::Ref data() = 0; + virtual Eigen::Block data() = 0; - virtual Eigen::Ref get() = 0; + virtual Eigen::Block get() = 0; /// \brief Name of the sensor in the telemetry. /// @@ -340,8 +340,8 @@ namespace jiminy const Eigen::VectorXd & a, const Eigen::VectorXd & uMotor, const ForceVector & fExternal) override final; - virtual Eigen::Ref get() override final; - virtual Eigen::Ref data() override final; + virtual Eigen::Block get() override final; + virtual Eigen::Block data() override final; private: virtual void attach(std::weak_ptr robot, diff --git a/core/include/jiminy/core/hardware/abstract_sensor.hxx b/core/include/jiminy/core/hardware/abstract_sensor.hxx index ab5d7479f..4e7ab6dda 100644 --- a/core/include/jiminy/core/hardware/abstract_sensor.hxx +++ b/core/include/jiminy/core/hardware/abstract_sensor.hxx @@ -289,14 +289,14 @@ namespace jiminy } template - inline Eigen::Ref AbstractSensorTpl::get() + inline Eigen::Block AbstractSensorTpl::get() { // No guard, since this method is not public return sharedStorage_->measurements_.col(sensorIndex_); } template - inline Eigen::Ref AbstractSensorTpl::data() + inline Eigen::Block AbstractSensorTpl::data() { // No guard, since this method is not public return sharedStorage_->data_.back().col(sensorIndex_); diff --git a/core/include/jiminy/core/hardware/basic_motors.h b/core/include/jiminy/core/hardware/basic_motors.h index 37b7330f7..cac400a60 100644 --- a/core/include/jiminy/core/hardware/basic_motors.h +++ b/core/include/jiminy/core/hardware/basic_motors.h @@ -17,6 +17,9 @@ namespace jiminy // Add extra options or update default values GenericConfig config = AbstractMotorBase::getDefaultMotorOptions(); + config["enableVelocityLimit"] = false; + config["velocityEffortInvSlope"] = 0.0; + config["enableEffortLimit"] = true; config["enableFriction"] = false; config["frictionViscousPositive"] = 0.0; config["frictionViscousNegative"] = 0.0; @@ -29,9 +32,15 @@ namespace jiminy struct SimpleMotorOptions : public AbstractMotorOptions { - /// \brief Flag to enable the joint friction. - /// - /// \pre Must be negative. + /// \brief Wether velocity limit is enabled. + const bool enableVelocityLimit; + /// \brief Inverse constant decrease rate of the maximum torque wrt velocity when + /// approaching the maximum velocity. The maximum torque is equal to zero at + /// maximum velocity. + const double velocityEffortInvSlope; + /// \brief Wether effort limit is enabled. + const bool enableEffortLimit; + /// \brief Wether joint friction is enabled. const bool enableFriction; /// \brief Viscous coefficient of the joint friction for positive velocity. /// @@ -58,6 +67,9 @@ namespace jiminy SimpleMotorOptions(const GenericConfig & options) : AbstractMotorOptions(options), + enableVelocityLimit(boost::get(options.at("enableVelocityLimit"))), + velocityEffortInvSlope{boost::get(options.at("velocityEffortInvSlope"))}, + enableEffortLimit(boost::get(options.at("enableEffortLimit"))), enableFriction{boost::get(options.at("enableFriction"))}, frictionViscousPositive{boost::get(options.at("frictionViscousPositive"))}, frictionViscousNegative{boost::get(options.at("frictionViscousNegative"))}, @@ -69,7 +81,7 @@ namespace jiminy }; public: - explicit SimpleMotor(const std::string & name) noexcept; + explicit SimpleMotor(const std::string & name); virtual ~SimpleMotor() = default; void initialize(const std::string & jointName); diff --git a/core/include/jiminy/core/hardware/basic_sensors.h b/core/include/jiminy/core/hardware/basic_sensors.h index 456fd3f42..b6f3d87d3 100644 --- a/core/include/jiminy/core/hardware/basic_sensors.h +++ b/core/include/jiminy/core/hardware/basic_sensors.h @@ -137,14 +137,15 @@ namespace jiminy public: using AbstractSensorTpl::AbstractSensorTpl; - void initialize(const std::string & jointName); + void initialize(const std::string & motorOrJointName, bool isJointSide = false); void setOptions(const GenericConfig & sensorOptions) override; void refreshProxies() override; + const std::string & getMotorName() const; + std::size_t getMotorIndex() const; const std::string & getJointName() const; pinocchio::JointIndex getJointIndex() const; - JointModelType getJointType() const; private: void set(double t, @@ -157,7 +158,14 @@ namespace jiminy private: std::string jointName_{}; pinocchio::JointIndex jointIndex_{0}; + std::string motorName_{}; + std::size_t motorIndex_{0}; + bool isJointSide_{false}; + JointModelType jointType_{JointModelType::UNSUPPORTED}; + std::size_t jointPositionIndex_{0}; + std::size_t joinVelocityIndex_{0}; + double mechanicalReduction_{1.0}; }; class EffortSensor; diff --git a/core/include/jiminy/core/io/serialization.h b/core/include/jiminy/core/io/serialization.h index ee847b559..860f77c1f 100644 --- a/core/include/jiminy/core/io/serialization.h +++ b/core/include/jiminy/core/io/serialization.h @@ -37,7 +37,7 @@ namespace jiminy template void loadFromBinary(T & obj, const std::string & data); - std::string saveToBinary(const std::shared_ptr & robot, + std::string saveToBinary(const std::shared_ptr & robot, bool isPersistent = true); void loadFromBinary(std::shared_ptr & robot, @@ -65,9 +65,16 @@ namespace boost::serialization // ***************************************** jiminy **************************************** // + template + void load_construct_data(Archive & ar, jiminy::Model * modelPtr, const unsigned int version); + template void serialize(Archive & ar, jiminy::Model & model, const unsigned int version); + template + void + save_construct_data(Archive & ar, const jiminy::Robot * robotPtr, const unsigned int version); + template void load_construct_data(Archive & ar, jiminy::Robot * robotPtr, const unsigned int version); diff --git a/core/include/jiminy/core/robot/model.h b/core/include/jiminy/core/robot/model.h index 7d05cf584..da35bd22b 100644 --- a/core/include/jiminy/core/robot/model.h +++ b/core/include/jiminy/core/robot/model.h @@ -139,9 +139,6 @@ namespace jiminy config["positionLimitFromUrdf"] = true; config["positionLimitMin"] = Eigen::VectorXd{}; config["positionLimitMax"] = Eigen::VectorXd{}; - config["enableVelocityLimit"] = true; - config["velocityLimitFromUrdf"] = true; - config["velocityLimit"] = Eigen::VectorXd{}; return config; }; @@ -186,17 +183,11 @@ namespace jiminy /// \brief Min position limit of all the mechanical joints of the theoretical model. const Eigen::VectorXd positionLimitMin; const Eigen::VectorXd positionLimitMax; - const bool enableVelocityLimit; - const bool velocityLimitFromUrdf; - const Eigen::VectorXd velocityLimit; JointOptions(const GenericConfig & options) : positionLimitFromUrdf{boost::get(options.at("positionLimitFromUrdf"))}, positionLimitMin{boost::get(options.at("positionLimitMin"))}, - positionLimitMax{boost::get(options.at("positionLimitMax"))}, - enableVelocityLimit{boost::get(options.at("enableVelocityLimit"))}, - velocityLimitFromUrdf{boost::get(options.at("velocityLimitFromUrdf"))}, - velocityLimit{boost::get(options.at("velocityLimit"))} + positionLimitMax{boost::get(options.at("positionLimitMax"))} { } }; @@ -247,9 +238,6 @@ namespace jiminy } }; - public: - JIMINY_DISABLE_COPY(Model) - public: /* Manually enforcing memory alignment. @@ -276,6 +264,9 @@ namespace jiminy public: explicit Model() noexcept; + Model(const Model & other); + Model & operator=(const Model & other); + Model & operator=(Model && other) = default; virtual ~Model() = default; void initialize( @@ -375,13 +366,10 @@ namespace jiminy const std::vector & getBacklashJointNames() const; const std::vector & getBacklashJointIndices() const; - const Eigen::VectorXd & getPositionLimitMin() const; - const Eigen::VectorXd & getPositionLimitMax() const; - const Eigen::VectorXd & getVelocityLimit() const; - const std::vector & getLogPositionFieldnames() const; const std::vector & getLogVelocityFieldnames() const; const std::vector & getLogAccelerationFieldnames() const; + const std::vector & getLogEffortFieldnames() const; const std::vector & getLogForceExternalFieldnames() const; void getExtendedPositionFromTheoretical(const Eigen::VectorXd & qTheoretical, @@ -477,20 +465,14 @@ namespace jiminy /// \brief Store constraints. ConstraintTree constraints_{}; - /// \brief Upper position limit of the whole configuration vector (INF for non-physical - /// joints, ie flexibility joints and freeflyer, if any). - Eigen::VectorXd positionLimitMin_{}; - /// \brief Lower position limit of the whole configuration vector. - Eigen::VectorXd positionLimitMax_{}; - /// \brief Maximum absolute velocity of the whole velocity vector. - Eigen::VectorXd velocityLimit_{}; - /// \brief Fieldnames of the elements in the configuration vector of the model. std::vector logPositionFieldnames_{}; /// \brief Fieldnames of the elements in the velocity vector of the model. std::vector logVelocityFieldnames_{}; /// \brief Fieldnames of the elements in the acceleration vector of the model. std::vector logAccelerationFieldnames_{}; + /// \brief Fieldnames of the elements in the effort vector of the model. + std::vector logEffortFieldnames_{}; /// \brief Concatenated fieldnames of the external force applied at each joint of the /// model, 'universe' excluded. std::vector logForceExternalFieldnames_{}; diff --git a/core/include/jiminy/core/robot/robot.h b/core/include/jiminy/core/robot/robot.h index b4a3f0572..ad4cae0d0 100644 --- a/core/include/jiminy/core/robot/robot.h +++ b/core/include/jiminy/core/robot/robot.h @@ -34,16 +34,19 @@ namespace jiminy }; public: + using WeakMotorVector = std::vector>; using MotorVector = std::vector>; + using WeakSensorVector = std::vector>; using SensorVector = std::vector>; + using WeakSensorTree = std::unordered_map; using SensorTree = std::unordered_map; - public: - JIMINY_DISABLE_COPY(Robot) - public: /// \param[in] name Name of the Robot. explicit Robot(const std::string & name = "") noexcept; + Robot(const Robot & other); + Robot & operator=(const Robot & other); + Robot & operator=(Robot && other) = default; virtual ~Robot(); auto shared_from_this() { return shared_from(this); } @@ -61,19 +64,22 @@ namespace jiminy const std::string & getName() const; void attachMotor(std::shared_ptr motor); - std::shared_ptr getMotor(const std::string & motorName); - std::weak_ptr getMotor(const std::string & motorName) const; - const MotorVector & getMotors() const; void detachMotor(const std::string & motorName); void detachMotors(std::vector motorNames = {}); void attachSensor(std::shared_ptr sensor); + void detachSensor(const std::string & sensorType, const std::string & sensorName); + void detachSensors(const std::string & sensorType = {}); + + std::shared_ptr getMotor(const std::string & motorName); + std::weak_ptr getMotor(const std::string & motorName) const; + const MotorVector & getMotors(); + const WeakMotorVector & getMotors() const; std::shared_ptr getSensor(const std::string & sensorType, const std::string & sensorName); std::weak_ptr getSensor(const std::string & sensorType, const std::string & sensorName) const; - const SensorTree & getSensors() const; - void detachSensor(const std::string & sensorType, const std::string & sensorName); - void detachSensors(const std::string & sensorType = {}); + const SensorTree & getSensors(); + const WeakSensorTree & getSensors() const; void setController(const std::shared_ptr & controller); std::shared_ptr getController(); @@ -84,8 +90,9 @@ namespace jiminy const Eigen::VectorXd & v, const Eigen::VectorXd & a, const Eigen::VectorXd & command); - const Eigen::VectorXd & getMotorEfforts() const; - double getMotorEffort(const std::string & motorName) const; + std::tuple, Eigen::Ref> + getMotorEfforts() const; + std::tuple getMotorEffort(const std::string & motorName) const; /// \warning It assumes that kinematic quantities have been updated previously and are /// consistent with the following input arguments. If not, one must call @@ -113,29 +120,15 @@ namespace jiminy void updateTelemetry(); bool getIsTelemetryConfigured() const; - const std::vector & getMotorNames() const; - std::vector getMotorJointIndices() const; - std::vector> getMotorsPositionIndices() const; - std::vector getMotorVelocityIndices() const; - const std::unordered_map> & getSensorNames() const; - const std::vector & getSensorNames(const std::string & sensorType) const; - - const Eigen::VectorXd & getCommandLimit() const; - const std::vector & getLogCommandFieldnames() const; - const std::vector & getLogMotorEffortFieldnames() const; - - // Getters without 'get' prefix for consistency with pinocchio C++ API - Eigen::Index nmotors() const; std::unique_ptr getLock(); bool getIsLocked() const; - protected: - void refreshMotorProxies(); - void refreshSensorProxies(); - void refreshProxies() override; + // Getters without 'get' prefix for consistency with pinocchio C++ API + Eigen::Index nmotors() const; + protected: void initializeExtendedModel() override; private: @@ -146,26 +139,26 @@ namespace jiminy std::shared_ptr telemetryData_{nullptr}; /// \brief Motors attached to the robot. MotorVector motors_{}; + /// \brief Motors attached to the robot as non-owning data structure of const references. + WeakMotorVector motorsWeakConstRef_{}; /// \brief Sensors attached to the robot. SensorTree sensors_{}; + /// \brief Sensors attached to the robot as non-owning data structure of const references. + WeakSensorTree sensorsWeakConstRef_{}; /// \brief Whether the robot options are guaranteed to be up-to-date. mutable bool areRobotOptionsRefreshed_{false}; /// \brief Dictionary with the parameters of the robot. mutable GenericConfig robotOptionsGeneric_{}; - /// \brief Name of the motors. - std::vector motorNames_{}; - /// \brief Name of the sensors. - std::unordered_map> sensorNames_{}; - /// \brief Fieldnames of the command. + /// \brief Log telemetry fieldnames associated with the command of each motor. std::vector logCommandFieldnames_{}; - /// \brief Fieldnames of the motors effort. - std::vector logMotorEffortFieldnames_{}; - /// \brief Number of motors. - Eigen::Index nmotors_{0}; /// \brief Controller of the robot. std::shared_ptr controller_{nullptr}; /// \brief Name of the robot. std::string name_; + /// \brief Velocity limit associated with each motor. + Eigen::VectorXd motorVelocityLimit_; + /// \brief Effort limit associated with each motor. + Eigen::VectorXd motorEffortLimit_; private: std::unique_ptr mutexLocal_{std::make_unique()}; diff --git a/core/include/jiminy/core/utilities/geometry.h b/core/include/jiminy/core/utilities/geometry.h index 37b201953..75ef89ca5 100644 --- a/core/include/jiminy/core/utilities/geometry.h +++ b/core/include/jiminy/core/utilities/geometry.h @@ -25,6 +25,29 @@ namespace jiminy double yMax, double yUnit, bool mustSimplify = false); + + /// \brief Unidirectional periodic stairs ground of even parity, consisting of alternating + /// ascending and descending staircases. + /// _______ _______ + /// __ _______|H |_______ _______|H |_______ + /// |______________|H H|_______|H H|_ . . . + /// . . . . . . . . . . + /// . W . W . W . W . W . W . W . W . W . + /// . i=0 . i=0 . i=1 . i=N . i=N+1 . i=0 . i=1 . i=N . i=N+1 . + /// |------> + /// x = 0 + /// + /// \details The stairs have identical height and width, and each staircase has an identical + /// step number. This number corresponds to the amount of steps to climb in order to + /// reach the highest steps from the lowest ones. The above ASCII art shows staircases + /// with a step number of two. + /// + /// \param[in] stepWidth Width of the steps. + /// \param[in] stepHeight Heigh of the steps. + /// \param[in] stepNumber Number of steps in the ascending or descending direction. + /// \param[in] orientation Orientation of the staircases in the XY plane. + HeightmapFunction JIMINY_DLLAPI stairs( + double stepWidth, double stepHeight, uint32_t stepNumber, double orientation); } #endif // JIMINY_GEOMETRY_H \ No newline at end of file diff --git a/core/src/constraints/frame_constraint.cc b/core/src/constraints/frame_constraint.cc index 5d99531c4..7512b76ce 100644 --- a/core/src/constraints/frame_constraint.cc +++ b/core/src/constraints/frame_constraint.cc @@ -135,8 +135,8 @@ namespace jiminy const pinocchio::SE3 & framePose = model->pinocchioData_.oMf[frameIndex_]; const pinocchio::SE3 transformLocal(rotationLocal_, framePose.translation()); const pinocchio::Frame & frame = model->pinocchioModel_.frames[frameIndex_]; - const pinocchio::JointModel & joint = model->pinocchioModel_.joints[frame.parent]; - const int32_t colRef = joint.nv() + joint.idx_v() - 1; + const pinocchio::JointModel & jmodel = model->pinocchioModel_.joints[frame.parent]; + const int32_t colRef = jmodel.nv() + jmodel.idx_v() - 1; for (Eigen::DenseIndex j = colRef; j >= 0; j = model->pinocchioData_.parents_fromRow[static_cast(j)]) { diff --git a/core/src/engine/engine.cc b/core/src/engine/engine.cc index 255cbf6be..a71b4943d 100644 --- a/core/src/engine/engine.cc +++ b/core/src/engine/engine.cc @@ -121,6 +121,7 @@ namespace jiminy command.setZero(nMotors); u.setZero(nv); uMotor.setZero(nMotors); + uTransmission.setZero(nMotors); uInternal.setZero(nv); uCustom.setZero(nv); fExternal = ForceVector(nJoints, pinocchio::Force::Zero()); @@ -140,6 +141,7 @@ namespace jiminy command.resize(0); u.resize(0); uMotor.resize(0); + uTransmission.resize(0); uInternal.resize(0); uCustom.resize(0); fExternal.clear(); @@ -599,6 +601,11 @@ namespace jiminy robotName, {}, TELEMETRY_FIELDNAME_DELIMITER); + robotDataIt->logEffortFieldnames = + addCircumfix((*robotIt)->getLogEffortFieldnames(), + robotName, + {}, + TELEMETRY_FIELDNAME_DELIMITER); robotDataIt->logForceExternalFieldnames = addCircumfix((*robotIt)->getLogForceExternalFieldnames(), robotName, @@ -609,11 +616,6 @@ namespace jiminy robotName, {}, TELEMETRY_FIELDNAME_DELIMITER); - robotDataIt->logMotorEffortFieldnames = - addCircumfix((*robotIt)->getLogMotorEffortFieldnames(), - robotName, - {}, - TELEMETRY_FIELDNAME_DELIMITER); robotDataIt->logEnergyFieldname = addCircumfix("energy", robotName, {}, TELEMETRY_FIELDNAME_DELIMITER); @@ -633,6 +635,11 @@ namespace jiminy telemetrySender_->registerVariable(robotDataIt->logAccelerationFieldnames, robotDataIt->state.a); } + if (engineOptions_->telemetry.enableEffort) + { + telemetrySender_->registerVariable(robotDataIt->logEffortFieldnames, + robotDataIt->state.u); + } if (engineOptions_->telemetry.enableForceExternal) { for (std::size_t i = 1; i < robotDataIt->state.fExternal.size(); ++i) @@ -651,11 +658,6 @@ namespace jiminy telemetrySender_->registerVariable(robotDataIt->logCommandFieldnames, robotDataIt->state.command); } - if (engineOptions_->telemetry.enableMotorEffort) - { - telemetrySender_->registerVariable(robotDataIt->logMotorEffortFieldnames, - robotDataIt->state.uMotor); - } if (engineOptions_->telemetry.enableEnergy) { telemetrySender_->registerVariable(robotDataIt->logEnergyFieldname, @@ -974,9 +976,8 @@ namespace jiminy /* Check that the initial configuration is not out-of-bounds if appropriate. Note that EPS allows to be very slightly out-of-bounds, which may occurs because of rounding errors. */ - if (((contactModel_ == ContactModelType::CONSTRAINT) && - ((EPS < q.array() - robot->getPositionLimitMax().array()).any() || - (EPS < robot->getPositionLimitMin().array() - q.array()).any()))) + if ((EPS < q.array() - robot->pinocchioModel_.upperPositionLimit.array()).any() || + (EPS < robot->pinocchioModel_.lowerPositionLimit.array() - q.array()).any()) { JIMINY_THROW(std::invalid_argument, "Initial configuration out-of-bounds for robot '", @@ -985,8 +986,7 @@ namespace jiminy } // Check that the initial velocity is not out-of-bounds - if ((robot->modelOptions_->joints.enableVelocityLimit && - (robot->getVelocityLimit().array() < v.array().abs()).any())) + if ((robot->pinocchioModel_.velocityLimit.array() < v.array().abs()).any()) { JIMINY_THROW(std::invalid_argument, "Initial velocity out-of-bounds for robot '", @@ -1371,6 +1371,7 @@ namespace jiminy Eigen::VectorXd & u = robotDataIt->state.u; Eigen::VectorXd & command = robotDataIt->state.command; Eigen::VectorXd & uMotor = robotDataIt->state.uMotor; + Eigen::VectorXd & uTransmission = robotDataIt->state.uTransmission; Eigen::VectorXd & uInternal = robotDataIt->state.uInternal; Eigen::VectorXd & uCustom = robotDataIt->state.uCustom; ForceVector & fext = robotDataIt->state.fExternal; @@ -1402,7 +1403,8 @@ namespace jiminy // Compute the actual motor effort (*robotIt)->computeMotorEfforts(t, q, v, a, command); - uMotor = (*robotIt)->getMotorEfforts(); + const auto & uMotorAndJoint = (*robotIt)->getMotorEfforts(); + std::tie(uMotor, uTransmission) = uMotorAndJoint; // Compute the internal dynamics uCustom.setZero(); @@ -1413,8 +1415,10 @@ namespace jiminy for (const auto & motor : (*robotIt)->getMotors()) { const std::size_t motorIndex = motor->getIndex(); - const Eigen::Index motorVelocityIndex = motor->getJointVelocityIndex(); - u[motorVelocityIndex] += uMotor[motorIndex]; + const pinocchio::JointIndex jointIndex = motor->getJointIndex(); + const Eigen::Index motorVelocityIndex = + (*robotIt)->pinocchioModel_.joints[jointIndex].idx_v(); + u[motorVelocityIndex] += uTransmission[motorIndex]; } } isFirstIter = false; @@ -2517,7 +2521,7 @@ namespace jiminy if (isSimulationRunning_) { JIMINY_THROW(bad_control_flow, - "Simulation already running. Stop it before removing coupling forces."); + "Simulation already running. Stop it before removing impulse forces."); } std::ptrdiff_t robotIndex = getRobotIndex(robotName); @@ -2531,7 +2535,7 @@ namespace jiminy if (isSimulationRunning_) { JIMINY_THROW(bad_control_flow, - "simulation already running. Stop it before removing coupling forces."); + "simulation already running. Stop it before removing impulse forces."); } for (auto & robotData : robotDataVec_) @@ -2546,7 +2550,7 @@ namespace jiminy if (isSimulationRunning_) { JIMINY_THROW(bad_control_flow, - "Simulation already running. Stop it before removing coupling forces."); + "Simulation already running. Stop it before removing profile forces."); } @@ -2569,7 +2573,7 @@ namespace jiminy if (isSimulationRunning_) { JIMINY_THROW(bad_control_flow, - "Simulation already running. Stop it before removing coupling forces."); + "Simulation already running. Stop it before removing profile forces."); } for (auto & robotData : robotDataVec_) @@ -3191,57 +3195,15 @@ namespace jiminy robot->getController()->computeCommand(t, q, v, command); } - template class JointModel, typename Scalar, int Options, int axis> - static std::enable_if_t< - is_pinocchio_joint_revolute_v> || - is_pinocchio_joint_revolute_unbounded_v>, - double> - getSubtreeInertiaProj(const JointModel & /* model */, - const pinocchio::Inertia & Isubtree) - { - double inertiaProj = Isubtree.inertia()(axis, axis); - for (Eigen::Index i = 0; i < 3; ++i) - { - if (i != axis) - { - inertiaProj += Isubtree.mass() * std::pow(Isubtree.lever()[i], 2); - } - } - return inertiaProj; - } - - template - static std::enable_if_t || - is_pinocchio_joint_revolute_unbounded_unaligned_v, - double> - getSubtreeInertiaProj(const JointModel & model, const pinocchio::Inertia & Isubtree) - { - return model.axis.dot(Isubtree.inertia() * model.axis) + - Isubtree.mass() * model.axis.cross(Isubtree.lever()).squaredNorm(); - } - - template - static std::enable_if_t || - is_pinocchio_joint_prismatic_unaligned_v, - double> - getSubtreeInertiaProj(const JointModel & /* model */, const pinocchio::Inertia & Isubtree) - { - return Isubtree.mass(); - } - struct computePositionLimitsForcesAlgo : public pinocchio::fusion::JointUnaryVisitorBase { typedef boost::fusion::vector< - const pinocchio::Data & /* pinocchioData */, const Eigen::VectorXd & /* q */, - const Eigen::VectorXd & /* v */, const Eigen::VectorXd & /* positionLimitMin */, const Eigen::VectorXd & /* positionLimitMax */, const std::unique_ptr & /* engineOptions */, - ContactModelType /* contactModel */, - const std::shared_ptr & /* constraint */, - Eigen::VectorXd & /* u */> + const std::shared_ptr & /* constraint */> ArgsType; template @@ -3251,64 +3213,33 @@ namespace jiminy is_pinocchio_joint_prismatic_unaligned_v, void> algo(const pinocchio::JointModelBase & joint, - const pinocchio::Data & data, const Eigen::VectorXd & q, - const Eigen::VectorXd & v, const Eigen::VectorXd & positionLimitMin, const Eigen::VectorXd & positionLimitMax, const std::unique_ptr & engineOptions, - ContactModelType contactModel, - const std::shared_ptr & constraint, - Eigen::VectorXd & u) + const std::shared_ptr & constraint) { // Define some proxies for convenience - const pinocchio::JointIndex jointIndex = joint.id(); const Eigen::Index positionIndex = joint.idx_q(); - const Eigen::Index velocityIndex = joint.idx_v(); const double qJoint = q[positionIndex]; const double qJointMin = positionLimitMin[positionIndex]; const double qJointMax = positionLimitMax[positionIndex]; - const double vJoint = v[velocityIndex]; - const double Ia = getSubtreeInertiaProj(joint.derived(), data.Ycrb[jointIndex]); - const double stiffness = engineOptions->joints.boundStiffness; - const double damping = engineOptions->joints.boundDamping; const double transitionEps = engineOptions->contacts.transitionEps; // Check if out-of-bounds - if (contactModel == ContactModelType::SPRING_DAMPER) + if (qJointMax < qJoint || qJoint < qJointMin) { - // Compute the acceleration to apply to move out of the bound - double accelJoint = 0.0; - if (qJoint > qJointMax) - { - const double qJointError = qJoint - qJointMax; - accelJoint = -std::max(stiffness * qJointError + damping * vJoint, 0.0); - } - else if (qJoint < qJointMin) - { - const double qJointError = qJoint - qJointMin; - accelJoint = -std::min(stiffness * qJointError + damping * vJoint, 0.0); - } - - // Apply the resulting force - u[velocityIndex] += Ia * accelJoint; + // Enable fixed joint constraint + auto & jointConstraint = static_cast(*constraint.get()); + jointConstraint.setReferenceConfiguration( + Eigen::Matrix(std::clamp(qJoint, qJointMin, qJointMax))); + jointConstraint.setRotationDir(qJointMax < qJoint); + constraint->enable(); } - else + else if (qJointMin + transitionEps < qJoint && qJoint < qJointMax - transitionEps) { - if (qJointMax < qJoint || qJoint < qJointMin) - { - // Enable fixed joint constraint - auto & jointConstraint = static_cast(*constraint.get()); - jointConstraint.setReferenceConfiguration( - Eigen::Matrix(std::clamp(qJoint, qJointMin, qJointMax))); - jointConstraint.setRotationDir(qJointMax < qJoint); - constraint->enable(); - } - else if (qJointMin + transitionEps < qJoint && qJoint < qJointMax - transitionEps) - { - // Disable fixed joint constraint - constraint->disable(); - } + // Disable fixed joint constraint + constraint->disable(); } } @@ -3317,21 +3248,14 @@ namespace jiminy is_pinocchio_joint_revolute_unbounded_unaligned_v, void> algo(const pinocchio::JointModelBase & /* joint */, - const pinocchio::Data & /* data */, const Eigen::VectorXd & /* q */, - const Eigen::VectorXd & /* v */, const Eigen::VectorXd & /* positionLimitMin */, const Eigen::VectorXd & /* positionLimitMax */, const std::unique_ptr & /* engineOptions */, - ContactModelType contactModel, - const std::shared_ptr & constraint, - Eigen::VectorXd & /* u */) + const std::shared_ptr & constraint) { - if (contactModel == ContactModelType::CONSTRAINT) - { - // Disable fixed joint constraint - constraint->disable(); - } + // Disable fixed joint constraint + constraint->disable(); } template @@ -3344,109 +3268,17 @@ namespace jiminy is_pinocchio_joint_composite_v, void> algo(const pinocchio::JointModelBase & /* joint */, - const pinocchio::Data & /* data */, const Eigen::VectorXd & /* q */, - const Eigen::VectorXd & /* v */, const Eigen::VectorXd & /* positionLimitMin */, const Eigen::VectorXd & /* positionLimitMax */, const std::unique_ptr & /* engineOptions */, - ContactModelType contactModel, - const std::shared_ptr & constraint, - Eigen::VectorXd & /* u */) + const std::shared_ptr & constraint) { #ifndef NDEBUG - JIMINY_WARNING("No position bounds implemented for this type of joint."); -#endif - if (contactModel == ContactModelType::CONSTRAINT) - { - // Disable fixed joint constraint - constraint->disable(); - } - } - }; - - struct computeVelocityLimitsForcesAlgo : - public pinocchio::fusion::JointUnaryVisitorBase - { - typedef boost::fusion::vector< - const pinocchio::Data & /* data */, - const Eigen::VectorXd & /* v */, - const Eigen::VectorXd & /* velocityLimitMax */, - const std::unique_ptr & /* engineOptions */, - ContactModelType /* contactModel */, - Eigen::VectorXd & /* u */> - ArgsType; - template - static std::enable_if_t || - is_pinocchio_joint_revolute_unaligned_v || - is_pinocchio_joint_revolute_unbounded_v || - is_pinocchio_joint_revolute_unbounded_unaligned_v || - is_pinocchio_joint_prismatic_v || - is_pinocchio_joint_prismatic_unaligned_v, - void> - algo(const pinocchio::JointModelBase & joint, - const pinocchio::Data & data, - const Eigen::VectorXd & v, - const Eigen::VectorXd & velocityLimitMax, - const std::unique_ptr & engineOptions, - ContactModelType contactModel, - Eigen::VectorXd & u) - { - // Define some proxies for convenience - const pinocchio::JointIndex jointIndex = joint.id(); - const Eigen::Index velocityIndex = joint.idx_v(); - const double vJoint = v[velocityIndex]; - const double vJointMin = -velocityLimitMax[velocityIndex]; - const double vJointMax = velocityLimitMax[velocityIndex]; - const double Ia = getSubtreeInertiaProj(joint.derived(), data.Ycrb[jointIndex]); - const double damping = engineOptions->joints.boundDamping; - - // Check if out-of-bounds - if (contactModel == ContactModelType::SPRING_DAMPER) - { - // Compute joint velocity error - double vJointError = 0.0; - if (vJoint > vJointMax) - { - vJointError = vJoint - vJointMax; - } - else if (vJoint < vJointMin) - { - vJointError = vJoint - vJointMin; - } - else - { - return; - } - - // Generate acceleration in the opposite direction if out-of-bounds - const double accelJoint = -2.0 * damping * vJointError; - - // Apply the resulting force - u[velocityIndex] += Ia * accelJoint; - } - } - - template - static std::enable_if_t || - is_pinocchio_joint_spherical_v || - is_pinocchio_joint_spherical_zyx_v || - is_pinocchio_joint_translation_v || - is_pinocchio_joint_planar_v || - is_pinocchio_joint_mimic_v || - is_pinocchio_joint_composite_v, - void> - algo(const pinocchio::JointModelBase & /* joint */, - const pinocchio::Data & /* data */, - const Eigen::VectorXd & /* v */, - const Eigen::VectorXd & /* velocityLimitMax */, - const std::unique_ptr & /* engineOptions */, - ContactModelType /* contactModel */, - Eigen::VectorXd & /* u */) - { -#ifndef NDEBUG - JIMINY_WARNING("No velocity bounds implemented for this type of joint."); + JIMINY_WARNING("Position bounds not implemented for this type of joint."); #endif + // Disable fixed joint constraint + constraint->disable(); } }; @@ -3458,13 +3290,12 @@ namespace jiminy { // Define some proxies const pinocchio::Model & model = robot->pinocchioModel_; - const pinocchio::Data & data = robot->pinocchioData_; const ConstraintTree & constraints = robot->getConstraints(); /* Enforce position limits for all joints having bounds constraints, ie mechanical and backlash joints. */ - const Eigen::VectorXd & positionLimitMin = robot->getPositionLimitMin(); - const Eigen::VectorXd & positionLimitMax = robot->getPositionLimitMax(); + const Eigen::VectorXd & positionLimitMin = robot->pinocchioModel_.lowerPositionLimit; + const Eigen::VectorXd & positionLimitMax = robot->pinocchioModel_.upperPositionLimit; for (auto & constraintItem : constraints.boundJoints) { auto & constraint = constraintItem.second; @@ -3472,34 +3303,11 @@ namespace jiminy const pinocchio::JointIndex jointIndex = jointConstraint->getJointIndex(); computePositionLimitsForcesAlgo::run( model.joints[jointIndex], - typename computePositionLimitsForcesAlgo::ArgsType(data, - q, - v, - positionLimitMin, - positionLimitMax, - engineOptions_, - contactModel_, - constraint, - uInternal)); - } - - // Enforce velocity limits for all joints having bounds constraints if requested - if (robot->modelOptions_->joints.enableVelocityLimit) - { - const Eigen::VectorXd & velocityLimitMax = robot->getVelocityLimit(); - for (auto & constraintItem : constraints.boundJoints) - { - auto & constraint = constraintItem.second; - const auto jointConstraint = std::static_pointer_cast(constraint); - const pinocchio::JointIndex jointIndex = jointConstraint->getJointIndex(); - computeVelocityLimitsForcesAlgo::run( - model.joints[jointIndex], - typename computeVelocityLimitsForcesAlgo::ArgsType( - data, v, velocityLimitMax, engineOptions_, contactModel_, uInternal)); - } + typename computePositionLimitsForcesAlgo::ArgsType( + q, positionLimitMin, positionLimitMax, engineOptions_, constraint)); } - // Compute the flexibilities (only support JointModelType::SPHERICAL so far) + // Compute the flexibilities (only support `JointModelType::SPHERICAL` so far) double angle; Eigen::Matrix3d rotJlog3; const Robot::DynamicsOptions & modelDynOptions = robot->modelOptions_->dynamics; @@ -3781,6 +3589,7 @@ namespace jiminy Eigen::VectorXd & u = robotDataIt->state.u; Eigen::VectorXd & command = robotDataIt->state.command; Eigen::VectorXd & uMotor = robotDataIt->state.uMotor; + Eigen::VectorXd & uTransmission = robotDataIt->state.uTransmission; Eigen::VectorXd & uInternal = robotDataIt->state.uInternal; Eigen::VectorXd & uCustom = robotDataIt->state.uCustom; ForceVector & fext = robotDataIt->state.fExternal; @@ -3817,7 +3626,8 @@ namespace jiminy /* Compute the actual motor effort. Note that it is impossible to have access to the current accelerations. */ (*robotIt)->computeMotorEfforts(t, *qIt, *vIt, aPrev, command); - uMotor = (*robotIt)->getMotorEfforts(); + const auto & uMotorAndJoint = (*robotIt)->getMotorEfforts(); + std::tie(uMotor, uTransmission) = uMotorAndJoint; /* Compute the user-defined internal dynamics. Make sure that the sensor state has been updated beforehand since the user-defined @@ -3830,8 +3640,10 @@ namespace jiminy for (const auto & motor : (*robotIt)->getMotors()) { const std::size_t motorIndex = motor->getIndex(); - const Eigen::Index motorVelocityIndex = motor->getJointVelocityIndex(); - u[motorVelocityIndex] += uMotor[motorIndex]; + const pinocchio::JointIndex jointIndex = motor->getJointIndex(); + const Eigen::Index motorVelocityIndex = + (*robotIt)->pinocchioModel_.joints[jointIndex].idx_v(); + u[motorVelocityIndex] += uTransmission[motorIndex]; } // Compute the dynamics diff --git a/core/src/hardware/abstract_motor.cc b/core/src/hardware/abstract_motor.cc index b8e4f880c..3331ddf13 100644 --- a/core/src/hardware/abstract_motor.cc +++ b/core/src/hardware/abstract_motor.cc @@ -6,9 +6,15 @@ namespace jiminy { - AbstractMotorBase::AbstractMotorBase(const std::string & name) noexcept : + AbstractMotorBase::AbstractMotorBase(const std::string & name) : name_{name} { + // Make sure that the motor name is not empty + if (name_.empty()) + { + JIMINY_THROW(std::logic_error, "Motor name must not be empty."); + } + // Initialize options motorOptionsGeneric_ = getDefaultMotorOptions(); setOptions(getOptions()); @@ -51,8 +57,8 @@ namespace jiminy motorIndex_ = sharedStorage_->num_; // Add a value for the motor to the shared data buffer - sharedStorage_->data_.conservativeResize(sharedStorage_->num_ + 1); - sharedStorage_->data_.tail<1>().setZero(); + sharedStorage_->data_.conservativeResize(sharedStorage_->num_ + 1, Eigen::NoChange); + sharedStorage_->data_.bottomRows<1>().setZero(); // Add the motor to the shared memory sharedStorage_->motors_.push_back(this); @@ -76,10 +82,10 @@ namespace jiminy { const Eigen::Index motorShift = static_cast(sharedStorage_->num_ - motorIndex_ - 1); - sharedStorage_->data_.segment(motorIndex_, motorShift) = - sharedStorage_->data_.tail(motorShift); + sharedStorage_->data_.middleRows(motorIndex_, motorShift) = + sharedStorage_->data_.bottomRows(motorShift); } - sharedStorage_->data_.conservativeResize(sharedStorage_->num_ - 1); + sharedStorage_->data_.conservativeResize(sharedStorage_->num_ - 1, Eigen::NoChange); // Shift the motor ids for (std::size_t i = motorIndex_ + 1; i < sharedStorage_->num_; ++i) @@ -163,18 +169,39 @@ namespace jiminy "Please stop it before setting motor options."); } - // Check if the internal buffers must be updated - if (isInitialized_) + /* Make sure that the mechanical reduction ratio is larger than 1.0. + While technically it is possible to have a reduction ratio smaller than 1.0, it is + extremely unlikely in practice as motors are supposed to spin fast to operate maximum + efficiency. It is explicitly forbidden here because it is error prone as the user may + assume the ratio to be inverted. */ + const double mechanicalReduction = + boost::get(motorOptions.at("mechanicalReduction")); + if (mechanicalReduction < 1.0 - EPS) { - // Check if armature has changed - const bool enableArmature = boost::get(motorOptions.at("enableArmature")); - mustNotifyRobot_ |= (baseMotorOptions_->enableArmature != enableArmature); - if (enableArmature) + JIMINY_THROW(std::invalid_argument, + "The mechanical reduction must be larger than 1.0."); + } + + /* Physically, it must be possible to deduce the joint position from the motor position. + For this condition to be satisfied, the joint position after doing exactly one turn on + motor side must remain the same (ignoring turns if any). This implies that the inverted + reduction ratio must be an integer. */ + if (jointType_ == JointModelType::ROTARY_UNBOUNDED) + { + if (abs(mechanicalReduction - 1.0) > EPS) { - const double armature = boost::get(motorOptions.at("armature")); - mustNotifyRobot_ |= // - std::abs(armature - baseMotorOptions_->armature) > EPS; + JIMINY_THROW(std::runtime_error, + "The mechanical reduction must be equal to 1.0 for motors actuating " + "unbounded revolute joints."); } + } + + // Check if the internal buffers must be updated + if (isInitialized_) + { + // Check if mechanical reduction ratio has changed + mustNotifyRobot_ |= abs(baseMotorOptions_->mechanicalReduction - mechanicalReduction) > + EPS; // Check if backlash has changed const bool enableBacklash = boost::get(motorOptions.at("enableBacklash")); @@ -186,14 +213,14 @@ namespace jiminy std::abs(backlash - baseMotorOptions_->backlash) > EPS; } - // Check if command limit has changed - const bool commandLimitFromUrdf = - boost::get(motorOptions.at("commandLimitFromUrdf")); - mustNotifyRobot_ |= (baseMotorOptions_->commandLimitFromUrdf != commandLimitFromUrdf); - if (!commandLimitFromUrdf) + // Check if armature has changed + const bool enableArmature = boost::get(motorOptions.at("enableArmature")); + mustNotifyRobot_ |= (baseMotorOptions_->enableArmature != enableArmature); + if (enableArmature) { - const double commandLimit = boost::get(motorOptions.at("commandLimit")); - mustNotifyRobot_ |= std::abs(commandLimit - baseMotorOptions_->commandLimit) > EPS; + const double armature = boost::get(motorOptions.at("armature")); + mustNotifyRobot_ |= // + std::abs(armature - baseMotorOptions_->armature) > EPS; } } @@ -203,13 +230,10 @@ namespace jiminy // Update inherited polymorphic accessor deepUpdate(motorOptionsGeneric_, motorOptions); - // Refresh the proxies if the robot is initialized if available - if (robot) + // Refresh proxies systematically if motor is attached + if (robot && isAttached_) { - if (mustNotifyRobot_ && robot->getIsInitialized() && isAttached_) - { - refreshProxies(); - } + refreshProxies(); } } @@ -220,12 +244,6 @@ namespace jiminy void AbstractMotorBase::refreshProxies() { - if (!isAttached_) - { - JIMINY_THROW(bad_control_flow, - "Motor not attached to any robot. Impossible to refresh motor proxies."); - } - auto robot = robot_.lock(); if (!robot) { @@ -233,6 +251,12 @@ namespace jiminy "Robot has been deleted. Impossible to refresh motor proxies."); } + if (!isAttached_) + { + JIMINY_THROW(bad_control_flow, + "Motor not attached to any robot. Impossible to refresh motor proxies."); + } + if (!isInitialized_) { JIMINY_THROW(bad_control_flow, @@ -245,8 +269,20 @@ namespace jiminy "Robot not initialized. Impossible to refresh motor proxies."); } - jointIndex_ = ::jiminy::getJointIndex(robot->pinocchioModel_, jointName_); - jointType_ = getJointTypeFromIndex(robot->pinocchioModel_, jointIndex_); + if (robot->getIsLocked()) + { + JIMINY_THROW(bad_control_flow, + "Robot already locked, probably because a simulation is running. " + "Please stop it before refreshing motor options."); + } + + // Define proxy for convenience + const double mechanicalReduction = baseMotorOptions_->mechanicalReduction; + const pinocchio::Model & model = robot->pinocchioModel_; + + // Get joint index and type + jointIndex_ = ::jiminy::getJointIndex(model, jointName_); + jointType_ = getJointTypeFromIndex(model, jointIndex_); // Motors are only supported for linear and rotary joints if (jointType_ != JointModelType::LINEAR && jointType_ != JointModelType::ROTARY && @@ -256,26 +292,51 @@ namespace jiminy "A motor can only be associated with a 1-dof linear or rotary joint."); } - jointPositionIndex_ = getJointPositionFirstIndex(robot->pinocchioModel_, jointName_); - jointVelocityIndex_ = getJointVelocityFirstIndex(robot->pinocchioModel_, jointName_); + // Deduce the motor position limits for the joint position limits + if (jointType_ == JointModelType::ROTARY_UNBOUNDED) + { + positionLimitLower_ = -INF; + positionLimitUpper_ = +INF; + } + else + { + const Eigen::Index jointPositionIndex = model.joints[jointIndex_].idx_q(); + positionLimitLower_ = + model.lowerPositionLimit[jointPositionIndex] * mechanicalReduction; + positionLimitUpper_ = + model.upperPositionLimit[jointPositionIndex] * mechanicalReduction; + } - // Get the motor effort limits from the URDF or the user options. - if (baseMotorOptions_->commandLimitFromUrdf) + // Get the motor effort limits on motor side from the URDF or the user options + if (baseMotorOptions_->effortLimitFromUrdf) { const Eigen::Index mechanicalJointVelocityIndex = getJointVelocityFirstIndex(robot->pinocchioModelTh_, jointName_); - commandLimit_ = robot->pinocchioModelTh_.effortLimit[mechanicalJointVelocityIndex] / - baseMotorOptions_->mechanicalReduction; + effortLimit_ = robot->pinocchioModelTh_.effortLimit[mechanicalJointVelocityIndex] / + mechanicalReduction; } else { - commandLimit_ = baseMotorOptions_->commandLimit; + effortLimit_ = baseMotorOptions_->effortLimit; } - // Get the rotor inertia + // Get the motor velocity limits on motor side from the URDF or the user options + if (baseMotorOptions_->velocityLimitFromUrdf) + { + const Eigen::Index mechanicalJointVelocityIndex = + getJointVelocityFirstIndex(robot->pinocchioModelTh_, jointName_); + velocityLimit_ = robot->pinocchioModelTh_.velocityLimit[mechanicalJointVelocityIndex] * + mechanicalReduction; + } + else + { + velocityLimit_ = baseMotorOptions_->velocityLimit; + } + + // Get the rotor inertia on joint side if (baseMotorOptions_->enableArmature) { - armature_ = baseMotorOptions_->armature; + armature_ = baseMotorOptions_->armature * std::pow(mechanicalReduction, 2); } else { @@ -301,24 +362,25 @@ namespace jiminy } } - double AbstractMotorBase::get() const + std::tuple AbstractMotorBase::get() const { - static double dataEmpty; if (isAttached_) { - return sharedStorage_->data_[motorIndex_]; + return {sharedStorage_->data_.coeff(motorIndex_, 0), + sharedStorage_->data_.coeff(motorIndex_, 1)}; } - return dataEmpty; + return {0.0, 0.0}; } - double & AbstractMotorBase::data() + std::tuple AbstractMotorBase::data() { - return sharedStorage_->data_[motorIndex_]; + return {sharedStorage_->data_(motorIndex_, 0), sharedStorage_->data_(motorIndex_, 1)}; } - const Eigen::VectorXd & AbstractMotorBase::getAll() const + std::tuple, Eigen::Ref> + AbstractMotorBase::getAll() const { - return sharedStorage_->data_; + return {sharedStorage_->data_.col(0), sharedStorage_->data_.col(1)}; } void AbstractMotorBase::setOptionsAll(const GenericConfig & motorOptions) @@ -365,24 +427,24 @@ namespace jiminy return jointIndex_; } - JointModelType AbstractMotorBase::getJointType() const + double AbstractMotorBase::getPositionLimitLower() const { - return jointType_; + return positionLimitLower_; } - Eigen::Index AbstractMotorBase::getJointPositionIndex() const + double AbstractMotorBase::getPositionLimitUpper() const { - return jointPositionIndex_; + return positionLimitUpper_; } - Eigen::Index AbstractMotorBase::getJointVelocityIndex() const + double AbstractMotorBase::getVelocityLimit() const { - return jointVelocityIndex_; + return velocityLimit_; } - double AbstractMotorBase::getCommandLimit() const + double AbstractMotorBase::getEffortLimit() const { - return commandLimit_; + return effortLimit_; } double AbstractMotorBase::getArmature() const @@ -407,22 +469,24 @@ namespace jiminy JIMINY_THROW(bad_control_flow, "Motor not attached to any robot."); } + // Make sure that the parent robot has not been deleted + auto robot = robot_.lock(); + if (!robot) + { + JIMINY_THROW(std::runtime_error, + "Robot has been deleted. Impossible to compute motor efforts."); + } + // Compute the actual effort of every motor for (AbstractMotorBase * motor : sharedStorage_->motors_) { - uint8_t nq_motor; - if (motor->getJointType() == JointModelType::ROTARY_UNBOUNDED) - { - nq_motor = 2; - } - else - { - nq_motor = 1; - } + const pinocchio::JointIndex jointIndex = motor->getJointIndex(); + const pinocchio::JointModel & jmodel = robot->pinocchioModel_.joints[jointIndex]; + const Eigen::Index jointVelocityIndex = jmodel.idx_v(); motor->computeEffort(t, - q.segment(motor->getJointPositionIndex(), nq_motor), - v[motor->getJointVelocityIndex()], - a[motor->getJointVelocityIndex()], + jmodel.jointConfigSelector(q), + v[jointVelocityIndex], + a[jointVelocityIndex], command[motor->getIndex()]); } } diff --git a/core/src/hardware/abstract_sensor.cc b/core/src/hardware/abstract_sensor.cc index 247aec850..569a57d59 100644 --- a/core/src/hardware/abstract_sensor.cc +++ b/core/src/hardware/abstract_sensor.cc @@ -6,11 +6,17 @@ namespace jiminy { - AbstractSensorBase::AbstractSensorBase(const std::string & name) noexcept : + AbstractSensorBase::AbstractSensorBase(const std::string & name) : generator_{std::seed_seq{std::random_device{}()}}, name_{name}, telemetrySender_{std::make_unique()} { + // Make sure that the motor name is not empty + if (name_.empty()) + { + JIMINY_THROW(std::logic_error, "Sensor name must not be empty."); + } + /* Initialize options. Note that the base implementation is called even if derived. This is a limitation of the C++ language specification which is not going to disappear anytime soon. */ diff --git a/core/src/hardware/basic_motors.cc b/core/src/hardware/basic_motors.cc index 448d9685d..b96d937ef 100644 --- a/core/src/hardware/basic_motors.cc +++ b/core/src/hardware/basic_motors.cc @@ -6,7 +6,7 @@ namespace jiminy { - SimpleMotor::SimpleMotor(const std::string & name) noexcept : + SimpleMotor::SimpleMotor(const std::string & name) : AbstractMotorBase(name) { /* AbstractMotorBase constructor calls the base implementations of the virtual methods @@ -18,21 +18,12 @@ namespace jiminy void SimpleMotor::initialize(const std::string & jointName) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - auto robot = robot_.lock(); - if (robot && robot->getIsLocked()) - { - JIMINY_THROW(bad_control_flow, - "Robot already locked, probably because a simulation is running. " - "Please stop it before refreshing motor proxies."); - } - // Update joint name jointName_ = jointName; isInitialized_ = true; - // Try refreshing proxies if possible, restore internals before throwing exception if not + /* Try refreshing proxies if possible, restore internals before throwing exception if not. + Note that it will raise an exception if a simulation is already running. */ try { refreshProxies(); @@ -70,6 +61,14 @@ namespace jiminy JIMINY_THROW(std::invalid_argument, "'frictionDrySlope' must be positive."); } + // Check that velocity limit is not enabled with effort limit + if (boost::get(motorOptions.at("enableVelocityLimit")) && + !boost::get(motorOptions.at("enableEffortLimit"))) + { + JIMINY_THROW(std::invalid_argument, + "'enableVelocityLimit' cannot be enabled without 'enableEffortLimit'."); + } + // Update class-specific "strongly typed" accessor for fast and convenient access motorOptions_ = std::make_unique(motorOptions); @@ -89,27 +88,47 @@ namespace jiminy "Motor not initialized. Impossible to compute actual motor effort."); } - /* Compute the motor effort, taking into account the limit, if any. + // Extract motor data + const auto & [uMotor, uTransmission] = data(); + + // Compute the velocity at motor-level + const double vMotor = motorOptions_->mechanicalReduction * v; + + /* Compute the motor effort, taking into account velocity and effort limits. It is the output of the motor on joint side, ie after the transmission. */ - if (motorOptions_->enableCommandLimit) + double effortMin = -INF; + double effortMax = INF; + if (motorOptions_->enableEffortLimit) { - command = std::clamp(command, -commandLimit_, commandLimit_); + effortMin = -effortLimit_; + effortMax = effortLimit_; + if (motorOptions_->enableVelocityLimit) + { + const double velocityThr = std::max( + velocityLimit_ - effortLimit_ * motorOptions_->velocityEffortInvSlope, 0.0); + effortMin *= std::clamp( + (velocityLimit_ + vMotor) / (velocityLimit_ - velocityThr), 0.0, 1.0); + effortMax *= std::clamp( + (velocityLimit_ - vMotor) / (velocityLimit_ - velocityThr), 0.0, 1.0); + } } - data() = motorOptions_->mechanicalReduction * command; + uMotor = std::clamp(command, effortMin, effortMax); + + // Translate motor effort on joint side + uTransmission = motorOptions_->mechanicalReduction * uMotor; - /* Add friction to the joints associated with the motor if enable. - It is computed on joint side instead of the motor. */ + // Add transmission friction if enable if (motorOptions_->enableFriction) { - if (v > 0) + if (v > 0.0) { - data() += + uTransmission += motorOptions_->frictionViscousPositive * v + motorOptions_->frictionDryPositive * tanh(motorOptions_->frictionDrySlope * v); } else { - data() += + uTransmission += motorOptions_->frictionViscousNegative * v + motorOptions_->frictionDryNegative * tanh(motorOptions_->frictionDrySlope * v); } diff --git a/core/src/hardware/basic_sensors.cc b/core/src/hardware/basic_sensors.cc index 8bbcc0d7a..8d6f24293 100644 --- a/core/src/hardware/basic_sensors.cc +++ b/core/src/hardware/basic_sensors.cc @@ -38,9 +38,15 @@ { \ JIMINY_THROW(bad_control_flow, \ "Sensor not initialized. Impossible to refresh sensor proxies."); \ + } \ + \ + if (robot->getIsLocked()) \ + { \ + JIMINY_THROW(bad_control_flow, \ + "Robot already locked, probably because a simulation is running. " \ + "Please stop it before refreshing sensor proxies."); \ } - #define GET_ROBOT_IF_INITIALIZED() \ if (!isInitialized_) \ { \ @@ -49,14 +55,6 @@ \ auto robot = robot_.lock(); -#define CHECK_SIMULATION_NOT_RUNNING() \ - auto robot = robot_.lock(); \ - if (robot && robot->getIsLocked()) \ - { \ - JIMINY_THROW(bad_control_flow, \ - "Robot already locked, probably because a simulation is running. " \ - "Please stop it before refreshing sensor proxies."); \ - } namespace jiminy { @@ -70,10 +68,6 @@ namespace jiminy void ImuSensor::initialize(const std::string & frameName) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - CHECK_SIMULATION_NOT_RUNNING() - // Update frame name frameName_ = frameName; isInitialized_ = true; @@ -114,7 +108,7 @@ namespace jiminy "gyroscope and accelerometer additive noise."); } - // Set options now that sanity check were made + // Set options now that sanity checks were made AbstractSensorTpl::setOptions(sensorOptions); } @@ -201,10 +195,6 @@ namespace jiminy void ContactSensor::initialize(const std::string & frameName) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - CHECK_SIMULATION_NOT_RUNNING() - // Update frame name frameName_ = frameName; isInitialized_ = true; @@ -237,7 +227,7 @@ namespace jiminy JIMINY_THROW(std::invalid_argument, "Wrong noise std vector size."); } - // Set options now that sanity check were made + // Set options now that sanity checks were made AbstractSensorTpl::setOptions(sensorOptions); } @@ -296,10 +286,6 @@ namespace jiminy void ForceSensor::initialize(const std::string & frameName) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - CHECK_SIMULATION_NOT_RUNNING() - // Update frame name frameName_ = frameName; isInitialized_ = true; @@ -332,7 +318,7 @@ namespace jiminy JIMINY_THROW(std::invalid_argument, "Wrong noise std vector size."); } - // Set options now that sanity check were made + // Set options now that sanity checks were made AbstractSensorTpl::setOptions(sensorOptions); } @@ -407,14 +393,18 @@ namespace jiminy template<> const std::vector AbstractSensorTpl::fieldnames_{"Q", "V"}; - void EncoderSensor::initialize(const std::string & jointName) + void EncoderSensor::initialize(const std::string & motorOrJointName, bool isJointSide) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - CHECK_SIMULATION_NOT_RUNNING() - - // Update joint name - jointName_ = jointName; + // Update motor name + isJointSide_ = isJointSide; + if (isJointSide_) + { + jointName_ = motorOrJointName; + } + else + { + motorName_ = motorOrJointName; + } isInitialized_ = true; // Try refreshing proxies if possible, restore internals before throwing exception if not @@ -424,7 +414,7 @@ namespace jiminy } catch (...) { - jointName_.clear(); + motorName_.clear(); isInitialized_ = false; throw; } @@ -445,7 +435,7 @@ namespace jiminy JIMINY_THROW(std::invalid_argument, "Wrong noise std vector size."); } - // Set options now that sanity check were made + // Set options now that sanity checks were made AbstractSensorTpl::setOptions(sensorOptions); } @@ -453,15 +443,25 @@ namespace jiminy { GET_ROBOT_AND_CHECK_SENSOR_INTEGRITY() - if (!robot->pinocchioModel_.existJointName(jointName_)) + // Refresh proxies that depends on whether the encoder is on the joint or motor side + if (isJointSide_) { - JIMINY_THROW(std::runtime_error, "Sensor attached to a joint that does not exist."); + motorIndex_ = robot->nmotors(); + jointIndex_ = ::jiminy::getJointIndex(robot->pinocchioModel_, jointName_); + mechanicalReduction_ = 1.0; + } + else + { + std::shared_ptr motor = robot->getMotor(motorName_).lock(); + motorIndex_ = motor->getIndex(); + jointIndex_ = motor->getJointIndex(); + jointName_ = robot->pinocchioModel_.names[jointIndex_]; + mechanicalReduction_ = motor->baseMotorOptions_->mechanicalReduction; } - - jointIndex_ = ::jiminy::getJointIndex(robot->pinocchioModel_, jointName_); - jointType_ = getJointTypeFromIndex(robot->pinocchioModel_, jointIndex_); // Motors are only supported for linear and rotary joints + const auto & jmodel = robot->pinocchioModel_.joints[jointIndex_]; + jointType_ = getJointType(jmodel); if (jointType_ != JointModelType::LINEAR && jointType_ != JointModelType::ROTARY && jointType_ != JointModelType::ROTARY_UNBOUNDED) { @@ -469,6 +469,20 @@ namespace jiminy std::runtime_error, "Encoder sensors can only be associated with a 1-dof linear or rotary joint."); } + + // Refresh joint position and velocity indices + jointPositionIndex_ = jmodel.idx_q(); + joinVelocityIndex_ = jmodel.idx_v(); + } + + const std::string & EncoderSensor::getMotorName() const + { + return motorName_; + } + + std::size_t EncoderSensor::getMotorIndex() const + { + return motorIndex_; } const std::string & EncoderSensor::getJointName() const @@ -481,11 +495,6 @@ namespace jiminy return jointIndex_; } - JointModelType EncoderSensor::getJointType() const - { - return jointType_; - } - void EncoderSensor::set(double /* t */, const Eigen::VectorXd & q, const Eigen::VectorXd & v, @@ -493,22 +502,29 @@ namespace jiminy const Eigen::VectorXd & /* uMotor */, const ForceVector & /* fExternal */) { - GET_ROBOT_IF_INITIALIZED() - - const auto & joint = robot->pinocchioModel_.joints[jointIndex_]; - const Eigen::Index jointPositionIndex = joint.idx_q(); - const Eigen::Index jointVelocityIndex = joint.idx_v(); + // Compute the joint position and velocity + double jointPosition; if (jointType_ == JointModelType::ROTARY_UNBOUNDED) { - const double cosTheta = q[jointPositionIndex]; - const double sinTheta = q[jointPositionIndex + 1]; - data()[0] = std::atan2(sinTheta, cosTheta); + const double cosTheta = q[jointPositionIndex_]; + const double sinTheta = q[jointPositionIndex_ + 1]; + jointPosition = std::atan2(sinTheta, cosTheta); + } + else + { + jointPosition = q[jointPositionIndex_]; + } + const double jointVelocity = v[joinVelocityIndex_]; + + // Compute encoder data depending on whether the encoder is on joint or motor side + if (isJointSide_) + { + data() << jointPosition, jointVelocity; } else { - data()[0] = q[jointPositionIndex]; + data() << jointPosition * mechanicalReduction_, jointVelocity * mechanicalReduction_; } - data()[1] = v[jointVelocityIndex]; } // ===================== EffortSensor ========================= @@ -520,10 +536,6 @@ namespace jiminy void EffortSensor::initialize(const std::string & motorName) { - // Make sure that no simulation is already running - // TODO: This check should be enforced by AbstractMotor somehow - CHECK_SIMULATION_NOT_RUNNING() - // Update motor name motorName_ = motorName; isInitialized_ = true; @@ -556,7 +568,7 @@ namespace jiminy JIMINY_THROW(std::invalid_argument, "Wrong noise std vector size."); } - // Set options now that sanity check were made + // Set options now that sanity checks were made AbstractSensorTpl::setOptions(sensorOptions); } diff --git a/core/src/io/json_writer.cc b/core/src/io/json_writer.cc index fce2f6fe9..0e4fbc2ed 100644 --- a/core/src/io/json_writer.cc +++ b/core/src/io/json_writer.cc @@ -28,9 +28,11 @@ namespace jiminy std::unique_ptr writer(builder.newStreamWriter()); std::ostream output(&buffer); writer->write(input, &output); - device_->resize(static_cast(buffer.str().size())); - device_->write(buffer.str()); + // FIXME: Use `view` to get a string_view rather than a string copy when moving to C++20 + const std::string data = buffer.str(); + device_->resize(static_cast(data.size())); + device_->write(data); device_->close(); } diff --git a/core/src/io/serialization.cc b/core/src/io/serialization.cc index d83045dc6..2f7bcfb56 100644 --- a/core/src/io/serialization.cc +++ b/core/src/io/serialization.cc @@ -32,6 +32,7 @@ #include "jiminy/core/hardware/abstract_sensor.h" #include "jiminy/core/hardware/basic_sensors.h" #include "jiminy/core/robot/robot.h" +#include "jiminy/core/utilities/pinocchio.h" #define BOOST_FILESYSTEM_VERSION 3 #include // `boost::filesystem::unique_path` @@ -640,7 +641,26 @@ namespace boost::serialization bool isPersistent = true; if constexpr (std::is_base_of_v) { - std::tie(isPersistent) = std::any_cast>(ar.state_); + if (ar.state_.has_value()) + { + try + { + std::tie(isPersistent) = std::any_cast>(ar.state_); + } + catch (const std::bad_any_cast & e) + { + JIMINY_WARNING("Failed to parse user-specified Model serialization arguments. " + "Using default values."); + } + } + } + + // Early return if not initalized + bool isInitialized = model.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; } // Backup URDF data @@ -653,8 +673,28 @@ namespace boost::serialization // Backup mesh package lookup directories ar << make_nvp("mesh_package_dirs", model.getMeshPackageDirs()); + /* Copy the theoretical model then remove extra collision frames. + Note that `removeCollisionBodies` is not called to avoid altering the robot. */ + pinocchio::Model pinocchioModelTh = model.pinocchioModelTh_; + std::vector collisionConstraintNames; + for (const std::string & name : model.getCollisionBodyNames()) + { + for (const pinocchio::GeometryObject & geom : model.collisionModelTh_.geometryObjects) + { + if (model.pinocchioModel_.frames[geom.parentFrame].name == name && + model.getConstraints().exist(geom.name, ConstraintNodeType::COLLISION_BODIES)) + { + const pinocchio::FrameIndex frameIndex = + getFrameIndex(pinocchioModelTh, geom.name); + pinocchioModelTh.frames.erase( + std::next(pinocchioModelTh.frames.begin(), frameIndex)); + pinocchioModelTh.nframes--; + } + } + } + // Backup theoretical and extended simulation models - ar << make_nvp("pinocchio_model_th", model.pinocchioModelTh_); + ar << make_nvp("pinocchio_model_th", pinocchioModelTh); ar << make_nvp("pinocchio_model", model.pinocchioModel_); /* Backup the Pinocchio GeometryModel for collisions and visuals if requested. @@ -739,8 +779,11 @@ namespace boost::serialization } template - void load(Archive & ar, Model & model, const unsigned int /* version */) + void load_construct_data(Archive & ar, Model * modelPtr, const unsigned int /* version */) { + // Create instance + ::new (modelPtr) Model(); + /* Tell the archive to start managing a shared_ptr. Note that this must be done manually here because, when a shared pointer is de-serialized, a raw pointer is first created. The latter only becomes managed by a @@ -761,19 +804,52 @@ namespace boost::serialization managed properly. This is a approach that has been preferred here. See: https://github.com/boostorg/serialization/blob/develop/include/boost/serialization/shared_ptr.hpp */ - std::shared_ptr modelPtr; + std::shared_ptr modelSharedPtr; shared_ptr_helper & h = ar.template get_helper>(shared_ptr_helper_id); - h.reset(modelPtr, &model); + h.reset(modelSharedPtr, modelPtr); + } + template + void load(Archive & ar, Model & model, const unsigned int /* version */) + { // Load arguments from archive state if available, otherwise use conservative default std::optional meshPathDir = std::nullopt; std::vector meshPackageDirs = {}; if constexpr (std::is_base_of_v) { - std::tie(meshPathDir, meshPackageDirs) = std::any_cast< - std::tuple &, const std::vector &>>( - ar.state_); + if (ar.state_.has_value()) + { + try + { + std::tie(meshPathDir, meshPackageDirs) = + std::any_cast &, + const std::vector &>>(ar.state_); + } + catch (const std::bad_any_cast & e) + { + JIMINY_WARNING("Failed to parse user-specified Model de-serialization " + "arguments. Using default values."); + } + } + } + + // Make sure that model is managed by a shared ptr + try + { + (void)model.shared_from_this(); + } + catch (std::bad_weak_ptr & e) + { + JIMINY_THROW(bad_control_flow, "Model must be managed by a std::shared_ptr."); + } + + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; } // Load URDF data @@ -892,8 +968,10 @@ namespace boost::serialization std::remove(urdfPath.c_str()); /* Restore the theoretical model. - Note that it is also restore manually-added frames since they are not tracked - internally for now. */ + Note that it also restores manually-added frames since they are not tracked + internally for now. Note that it is also necessary to restore the extended + simulation model otherwise adding collision bodies and contact points would fail. */ + model.pinocchioModel_ = pinocchioModelTh; model.pinocchioModelTh_ = pinocchioModelTh; model.pinocchioDataTh_ = pinocchio::Data(pinocchioModelTh); } @@ -937,6 +1015,7 @@ namespace boost::serialization } EXPLICIT_TEMPLATE_INSTANTIATION_SERIALIZE(Model) + EXPLICIT_TEMPLATE_INSTANTIATION_LOAD_CONSTRUCT(Model) } BOOST_CLASS_EXPORT(Model) @@ -974,6 +1053,14 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_motor_base", base_object(motor)); + // Early return if not initalized + bool isInitialized = motor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Save initialization data ar << make_nvp("joint_name", motor.getJointName()); } @@ -993,7 +1080,14 @@ namespace boost::serialization template void load(Archive & ar, SimpleMotor & motor, const unsigned int /* version */) { - // Tell the archive to start managing a shared_ptr + // Tell the archive to start managing a shared_ptr. + // Note that it is tricky to expose `load` to the user because it would require having + // access to the shared pointer managing the motor, along with the robot to which the motor + // should be attached. The only way to specify them is through the archive state. + // To get around this issue, we make sure that motors can only be serialized/de-serialized + // through the robot to which they are attached. This way, memory is guaranteed to be + // heap-allocated using `new` operator, and therefore requesting the archive to manage a + // shared_ptr of the motor instance is safe. */ std::shared_ptr motorPtr; shared_ptr_helper & h = ar.template get_helper>(shared_ptr_helper_id); @@ -1009,6 +1103,14 @@ namespace boost::serialization robot->attachMotor(motorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize motor only if attached std::string jointName; ar >> make_nvp("joint_name", jointName); @@ -1082,6 +1184,14 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_sensor_base", base_object(sensor)); + // Early return if not initalized + bool isInitialized = sensor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + /* Save initialization data. TODO: maybe `sensorRotationBiasInv_` should also be restored. */ ar << make_nvp("frame_name", sensor.getFrameName()); @@ -1117,6 +1227,14 @@ namespace boost::serialization robot->attachSensor(sensorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize sensor only if attached std::string frameName; ar >> make_nvp("frame_name", frameName); @@ -1156,6 +1274,14 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_sensor_base", base_object(sensor)); + // Early return if not initalized + bool isInitialized = sensor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Save initialization data ar << make_nvp("frame_name", sensor.getFrameName()); } @@ -1191,6 +1317,14 @@ namespace boost::serialization robot->attachSensor(sensorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize sensor only if attached std::string frameName; ar >> make_nvp("frame_name", frameName); @@ -1230,6 +1364,14 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_sensor_base", base_object(sensor)); + // Early return if not initalized + bool isInitialized = sensor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Save initialization data ar << make_nvp("frame_name", sensor.getFrameName()); } @@ -1265,6 +1407,14 @@ namespace boost::serialization robot->attachSensor(sensorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize sensor only if attached std::string frameName; ar >> make_nvp("frame_name", frameName); @@ -1286,7 +1436,7 @@ namespace boost::serialization BOOST_CLASS_EXPORT(ForceSensor) -// *********************************** jiminy::EncoderSensor *********************************** // +// ******************************** jiminy::EncoderSensor ********************************* // namespace boost::serialization { @@ -1304,8 +1454,25 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_sensor_base", base_object(sensor)); + // Early return if not initalized + bool isInitialized = sensor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Save initialization data - ar << make_nvp("joint_name", sensor.getJointName()); + const bool isJointSide = sensor.getMotorName().empty(); + ar << make_nvp("is_joint_side", isJointSide); + if (isJointSide) + { + ar << make_nvp("joint_name", sensor.getJointName()); + } + else + { + ar << make_nvp("motor_name", sensor.getMotorName()); + } } template @@ -1339,12 +1506,29 @@ namespace boost::serialization robot->attachSensor(sensorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize sensor only if attached - std::string jointName; - ar >> make_nvp("joint_name", jointName); + bool isJointSide; + ar >> make_nvp("is_joint_side", isJointSide); + std::string motorOrJointName; + if (isJointSide) + { + ar >> make_nvp("joint_name", motorOrJointName); + } + else + { + ar >> make_nvp("motor_name", motorOrJointName); + } if (sensor.getIsAttached()) { - sensor.initialize(jointName); + sensor.initialize(motorOrJointName, isJointSide); } } @@ -1378,6 +1562,14 @@ namespace boost::serialization // Save base ar << make_nvp("abstract_sensor_base", base_object(sensor)); + // Early return if not initalized + bool isInitialized = sensor.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Save initialization data ar << make_nvp("motor_name", sensor.getMotorName()); } @@ -1413,6 +1605,14 @@ namespace boost::serialization robot->attachSensor(sensorPtr); } + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Load initialization data but initialize sensor only if attached std::string motorName; ar >> make_nvp("motor_name", motorName); @@ -1452,24 +1652,32 @@ namespace boost::serialization // Save base model ar << make_nvp("model", base_object(robot)); + // Early return if not initalized + bool isInitialized = robot.getIsInitialized(); + ar << make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Backup all the motors using failsafe approach const std::shared_ptr dummyMotorPtr{}; - const Robot::MotorVector & motors = robot.getMotors(); - const std::vector & motorNames = robot.getMotorNames(); + const Robot::WeakMotorVector & motors = robot.getMotors(); std::size_t nmotors = robot.nmotors(); ar << make_nvp("nmotors", nmotors); for (std::size_t i = 0; i < nmotors; ++i) { + auto motor = std::const_pointer_cast(motors[i].lock()); const std::string name = toString("motor_", i); try { - ar << make_nvp(name.c_str(), motors[i]); + ar << make_nvp(name.c_str(), motor); } catch (const boost::archive::archive_exception & e) { ar << make_nvp(name.c_str(), dummyMotorPtr); JIMINY_WARNING("Failed to serialize motor '", - motorNames[i], + motor->getName(), "'. It will be missing when loading the robot from log." "\nRaised from exception: ", e.what()); @@ -1478,14 +1686,11 @@ namespace boost::serialization // Backup all the sensors using failsafe approach const std::shared_ptr dummySensorPtr{}; - const Robot::SensorTree & sensors = robot.getSensors(); - const std::unordered_map> & sensorNames = - robot.getSensorNames(); - const std::size_t nSensorTypes = sensorNames.size(); + const Robot::WeakSensorTree & sensors = robot.getSensors(); + const std::size_t nSensorTypes = sensors.size(); ar << make_nvp("nsensortypes", nSensorTypes); auto sensorsGroupIt = sensors.cbegin(); - auto sensorNamesGroupIt = sensorNames.cbegin(); - for (std::size_t i = 0; i < nSensorTypes; ++sensorsGroupIt, ++sensorNamesGroupIt, ++i) + for (std::size_t i = 0; i < nSensorTypes; ++sensorsGroupIt, ++i) { bool hasFailed = false; const auto & [sensorsGroupName, sensorsGroup] = *sensorsGroupIt; @@ -1493,12 +1698,13 @@ namespace boost::serialization ar << make_nvp(toString("type_", i, "_nsensors").c_str(), nSensors); for (std::size_t j = 0; j < nSensors; ++j) { + auto sensor = std::const_pointer_cast(sensorsGroup[j].lock()); const std::string name = toString("type_", i, "_sensor_", j); try { if (!hasFailed) { - ar << make_nvp(name.c_str(), sensorsGroup[j]); + ar << make_nvp(name.c_str(), sensor); } } catch (const boost::archive::archive_exception & e) @@ -1533,6 +1739,12 @@ namespace boost::serialization // Create instance ::new (robotPtr) Robot(name); + + // Tell the archive to start managing a shared_ptr + std::shared_ptr robotSharedPtr; + shared_ptr_helper & h = + ar.template get_helper>(shared_ptr_helper_id); + h.reset(robotSharedPtr, robotPtr); } template @@ -1541,6 +1753,14 @@ namespace boost::serialization // Load base model ar >> make_nvp("model", base_object(robot)); + // Early return if not initalized + bool isInitialized; + ar >> make_nvp("is_initialized", isInitialized); + if (!isInitialized) + { + return; + } + // Backup the already restored extended simulation model const pinocchio::Model pinocchioModel = robot.pinocchioModel_; @@ -1581,17 +1801,15 @@ namespace boost::serialization GenericConfig robotOptions; ar >> make_nvp("robot_options", robotOptions); GenericConfig & motorsOptions = boost::get(robotOptions.at("motors")); - const std::vector & motorNames = robot.getMotorNames(); - auto removeMotorOptionsPred = [&motorNames](const auto & motorOptionsItem) - { - const std::string & motorName = motorOptionsItem.first; - return std::find(motorNames.begin(), motorNames.end(), motorName) == motorNames.end(); - }; // FIXME: Replace by `std::erase_if` when moving to C++20 + const Robot::MotorVector & motors = robot.getMotors(); auto motorsOptionsIt = motorsOptions.begin(); while (motorsOptionsIt != motorsOptions.end()) { - if (removeMotorOptionsPred(*motorsOptionsIt)) + if (std::find_if(motors.begin(), + motors.end(), + [&motorName = motorsOptionsIt->first](const auto & motor) + { return motor->getName() == motorName; }) == motors.end()) { motorsOptionsIt = motorsOptions.erase(motorsOptionsIt); } @@ -1601,18 +1819,12 @@ namespace boost::serialization } } GenericConfig & sensorsOptions = boost::get(robotOptions.at("sensors")); - const std::unordered_map> & sensorNames = - robot.getSensorNames(); - auto removeSensorOptionsPred = [&sensorNames](const auto & sensorsGroupOptionsItem) - { - const std::string & sensorType = sensorsGroupOptionsItem.first; - return sensorNames.find(sensorType) == sensorNames.end(); - }; // FIXME: Replace by `std::erase_if` when moving to C++20 + const Robot::SensorTree & sensors = robot.getSensors(); auto sensorsOptionsIt = sensorsOptions.begin(); while (sensorsOptionsIt != sensorsOptions.end()) { - if (removeSensorOptionsPred(*sensorsOptionsIt)) + if (sensors.find(sensorsOptionsIt->first) == sensors.end()) { sensorsOptionsIt = sensorsOptions.erase(sensorsOptionsIt); } @@ -1641,7 +1853,7 @@ BOOST_CLASS_EXPORT(Robot) namespace jiminy { - std::string saveToBinary(const std::shared_ptr & robot, bool isPersistent) + std::string saveToBinary(const std::shared_ptr & robot, bool isPersistent) { return saveToBinaryImpl(robot, isPersistent); } diff --git a/core/src/robot/model.cc b/core/src/robot/model.cc index 9b410c2f5..36de4b91a 100644 --- a/core/src/robot/model.cc +++ b/core/src/robot/model.cc @@ -22,14 +22,17 @@ #include +#include + #include "urdf_parser/urdf_parser.h" -#include "jiminy/core/hardware/basic_sensors.h" #include "jiminy/core/robot/pinocchio_overload_algorithms.h" #include "jiminy/core/constraints/abstract_constraint.h" #include "jiminy/core/constraints/joint_constraint.h" #include "jiminy/core/constraints/sphere_constraint.h" #include "jiminy/core/constraints/frame_constraint.h" +#include "jiminy/core/hardware/basic_sensors.h" +#include "jiminy/core/io/serialization.h" #include "jiminy/core/utilities/pinocchio.h" #include "jiminy/core/utilities/helpers.h" @@ -193,6 +196,33 @@ namespace jiminy setOptions(getOptions()); } + Model::Model(const Model & other) : + enable_shared_from_this() + { + *this = other; + } + + Model & Model::operator=(const Model & other) + { + // Serialize the user-specified robot + std::stringstream sstream; + { + stateful_binary_oarchive oa(sstream); + const Model * const otherPtr = &other; + oa << boost::serialization::make_nvp("px", otherPtr); + } + + // De-serialize the user-specified robot then move assign to this robot + { + stateful_binary_iarchive ia(sstream); + Model * robotPtr; + ia >> boost::serialization::make_nvp("px", robotPtr); + *this = std::move(*robotPtr); + } + + return *this; + } + void initializePinocchioData(const pinocchio::Model & model, pinocchio::Data & data) { // Re-allocate Pinocchio Data from scratch @@ -215,7 +245,7 @@ namespace jiminy { (void)shared_from_this(); } - catch (std::bad_weak_ptr & e) + catch (std::bad_weak_ptr &) { JIMINY_THROW(bad_control_flow, "Model must be managed by a std::shared_ptr."); } @@ -505,7 +535,7 @@ namespace jiminy // Make sure that all the bodies exist for (const std::string & name : bodyNames) { - if (!pinocchioModel_.existBodyName(name)) + if (!pinocchioModelTh_.existBodyName(name)) { JIMINY_THROW(std::invalid_argument, "At least one of the bodies does not exist."); } @@ -685,7 +715,7 @@ namespace jiminy // Make sure that all the frames exist for (const std::string & name : frameNames) { - if (!pinocchioModel_.existFrame(name)) + if (!pinocchioModelTh_.existFrame(name)) { JIMINY_THROW(std::invalid_argument, "At least one of the frames does not exist."); } @@ -882,6 +912,23 @@ namespace jiminy return constraints_.exist(constraintName); } + bool Model::hasConstraints() const + { + bool hasConstraintsEnabled = false; + const_cast(constraints_) + .foreach( + [&hasConstraintsEnabled]( + const std::shared_ptr & constraint, + ConstraintNodeType /* node */) + { + if (constraint->getIsEnabled()) + { + hasConstraintsEnabled = true; + } + }); + return hasConstraintsEnabled; + } + void Model::resetConstraints(const Eigen::VectorXd & q, const Eigen::VectorXd & v) { constraints_.foreach([&q, &v](const std::shared_ptr & constraint, @@ -1167,6 +1214,8 @@ namespace jiminy logVelocityFieldnames_.reserve(static_cast(nv_)); logAccelerationFieldnames_.clear(); logAccelerationFieldnames_.reserve(static_cast(nv_)); + logEffortFieldnames_.clear(); + logEffortFieldnames_.reserve(static_cast(nv_)); logForceExternalFieldnames_.clear(); logForceExternalFieldnames_.reserve(6U * (pinocchioModel_.njoints - 1)); for (std::size_t i = 1; i < pinocchioModel_.joints.size(); ++i) @@ -1203,10 +1252,11 @@ namespace jiminy toString(jointPrefix, "Velocity", jointShortName, suffix)); logAccelerationFieldnames_.emplace_back( toString(jointPrefix, "Acceleration", jointShortName, suffix)); + logEffortFieldnames_.emplace_back( + toString(jointPrefix, "Effort", jointShortName, suffix)); } // Define complete external force fieldnames and backup them - std::vector jointForceExternalFieldnames; for (const std::string & suffix : ForceSensor::fieldnames_) { logForceExternalFieldnames_.emplace_back( @@ -1214,39 +1264,9 @@ namespace jiminy } } - // Get mechanical joint position limits from the URDF or user options - positionLimitMin_.setConstant(pinocchioModel_.nq, -INF); - positionLimitMax_.setConstant(pinocchioModel_.nq, +INF); - if (modelOptions_->joints.positionLimitFromUrdf) - { - for (Eigen::Index positionIndex : mechanicalJointPositionIndices_) - { - positionLimitMin_[positionIndex] = - pinocchioModel_.lowerPositionLimit[positionIndex]; - positionLimitMax_[positionIndex] = - pinocchioModel_.upperPositionLimit[positionIndex]; - } - } - else - { - for (std::size_t i = 0; i < mechanicalJointPositionIndices_.size(); ++i) - { - Eigen::Index positionIndex = mechanicalJointPositionIndices_[i]; - positionLimitMin_[positionIndex] = modelOptions_->joints.positionLimitMin[i]; - positionLimitMax_[positionIndex] = modelOptions_->joints.positionLimitMax[i]; - } - } - - // Get the backlash joint position limits - for (pinocchio::JointIndex jointIndex : backlashJointIndices_) - { - const Eigen::Index positionIndex = pinocchioModel_.idx_qs[jointIndex]; - positionLimitMin_[positionIndex] = pinocchioModel_.lowerPositionLimit[positionIndex]; - positionLimitMax_[positionIndex] = pinocchioModel_.upperPositionLimit[positionIndex]; - } - - /* Overwrite the position bounds for some specific joint type, mainly due to quaternion - normalization and cos/sin representation. */ + // Re-initialize position limits + Eigen::VectorXd positionLimitMin = Eigen::VectorXd::Constant(pinocchioModel_.nq, -INF); + Eigen::VectorXd positionLimitMax = Eigen::VectorXd::Constant(pinocchioModel_.nq, +INF); Eigen::Index idx_q, nq; for (const auto & joint : pinocchioModel_.joints) { @@ -1269,31 +1289,53 @@ namespace jiminy default: continue; } - positionLimitMin_.segment(idx_q, nq).setConstant(-1.0 - EPS); - positionLimitMax_.segment(idx_q, nq).setConstant(+1.0 + EPS); + positionLimitMin.segment(idx_q, nq).setConstant(-1.0 - EPS); + positionLimitMax.segment(idx_q, nq).setConstant(+1.0 + EPS); } - // Get the joint velocity limits from the URDF or the user options - velocityLimit_.setConstant(pinocchioModel_.nv, +INF); - if (modelOptions_->joints.enableVelocityLimit) + // Copy backlash joint position limits as-is from extended model + for (pinocchio::JointIndex jointIndex : backlashJointIndices_) + { + const Eigen::Index positionIndex = pinocchioModel_.idx_qs[jointIndex]; + positionLimitMin[positionIndex] = pinocchioModel_.lowerPositionLimit[positionIndex]; + positionLimitMax[positionIndex] = pinocchioModel_.upperPositionLimit[positionIndex]; + } + + // Set mechanical joint position limits from theoretical model or user options + if (modelOptions_->joints.positionLimitFromUrdf) { - if (modelOptions_->joints.velocityLimitFromUrdf) + int jointIndexTh = 1; + for (std::size_t i = 0; i < mechanicalJointIndices_.size(); ++i) { - for (Eigen::Index & velocityIndex : mechanicalJointVelocityIndices_) + const pinocchio::JointIndex jointIndex = mechanicalJointIndices_[i]; + const std::string & jointName = pinocchioModel_.names[jointIndex]; + while (pinocchioModelTh_.names[jointIndexTh] != jointName) { - velocityLimit_[velocityIndex] = pinocchioModel_.velocityLimit[velocityIndex]; + ++jointIndexTh; } + const auto & joint = pinocchioModel_.joints[jointIndex]; + const auto & jointTh = pinocchioModelTh_.joints[jointIndexTh]; + positionLimitMin.segment(joint.idx_q(), joint.nq()) = + pinocchioModelTh_.lowerPositionLimit.segment(jointTh.idx_q(), jointTh.nq()); + positionLimitMax.segment(joint.idx_q(), joint.nq()) = + pinocchioModelTh_.upperPositionLimit.segment(jointTh.idx_q(), jointTh.nq()); } - else + } + else + { + for (std::size_t i = 0; i < mechanicalJointPositionIndices_.size(); ++i) { - for (std::size_t i = 0; i < mechanicalJointVelocityIndices_.size(); ++i) - { - Eigen::Index velocityIndex = mechanicalJointVelocityIndices_[i]; - velocityLimit_[velocityIndex] = modelOptions_->joints.velocityLimit[i]; - } + Eigen::Index positionIndex = mechanicalJointPositionIndices_[i]; + positionLimitMin[positionIndex] = modelOptions_->joints.positionLimitMin[i]; + positionLimitMax[positionIndex] = modelOptions_->joints.positionLimitMax[i]; } } + // Overwrite extended model position limits + pinocchioModel_.lowerPositionLimit = positionLimitMin; + pinocchioModel_.upperPositionLimit = positionLimitMax; + + // Refresh all other proxies refreshGeometryProxies(); refreshContactProxies(); refreshConstraintProxies(); @@ -1468,30 +1510,6 @@ namespace jiminy internalBuffersMustBeUpdated = true; } } - bool velocityLimitFromUrdf = - boost::get(jointOptionsHolder.at("velocityLimitFromUrdf")); - if (!velocityLimitFromUrdf) - { - const Eigen::VectorXd & jointsVelocityLimit = - boost::get(jointOptionsHolder.at("velocityLimit")); - if (mechanicalJointVelocityIndices_.size() != - static_cast(jointsVelocityLimit.size())) - { - JIMINY_THROW(std::invalid_argument, "Wrong vector size for 'velocityLimit'."); - } - if (mechanicalJointVelocityIndices_.size() == - static_cast(modelOptions_->joints.velocityLimit.size())) - { - auto jointsVelocityLimitDiff = - jointsVelocityLimit - modelOptions_->joints.velocityLimit; - internalBuffersMustBeUpdated |= - (jointsVelocityLimitDiff.array().abs() >= EPS).all(); - } - else - { - internalBuffersMustBeUpdated = true; - } - } // Check if deformation points are all associated with different joints/frames const GenericConfig & dynOptionsHolder = @@ -1528,23 +1546,6 @@ namespace jiminy } } - // Check if the position or velocity limits have changed, and refresh proxies if so - bool enableVelocityLimit = - boost::get(jointOptionsHolder.at("enableVelocityLimit")); - if (positionLimitFromUrdf != modelOptions_->joints.positionLimitFromUrdf) - { - internalBuffersMustBeUpdated = true; - } - else if (enableVelocityLimit != modelOptions_->joints.enableVelocityLimit) - { - internalBuffersMustBeUpdated = true; - } - else if (enableVelocityLimit && - (velocityLimitFromUrdf != modelOptions_->joints.velocityLimitFromUrdf)) - { - internalBuffersMustBeUpdated = true; - } - // Check if the extended model must be regenerated bool enableFlexibility = boost::get(dynOptionsHolder.at("enableFlexibility")); if (modelOptions_ && @@ -1790,6 +1791,21 @@ namespace jiminy } } + Eigen::Index Model::nq() const + { + return nq_; + } + + Eigen::Index Model::nv() const + { + return nv_; + } + + Eigen::Index Model::nx() const + { + return nx_; + } + const std::vector & Model::getCollisionBodyNames() const { return collisionBodyNames_; @@ -1815,41 +1831,6 @@ namespace jiminy return contactFrameIndices_; } - const std::vector & Model::getLogPositionFieldnames() const - { - return logPositionFieldnames_; - } - - const Eigen::VectorXd & Model::getPositionLimitMin() const - { - return positionLimitMin_; - } - - const Eigen::VectorXd & Model::getPositionLimitMax() const - { - return positionLimitMax_; - } - - const std::vector & Model::getLogVelocityFieldnames() const - { - return logVelocityFieldnames_; - } - - const Eigen::VectorXd & Model::getVelocityLimit() const - { - return velocityLimit_; - } - - const std::vector & Model::getLogAccelerationFieldnames() const - { - return logAccelerationFieldnames_; - } - - const std::vector & Model::getLogForceExternalFieldnames() const - { - return logForceExternalFieldnames_; - } - const std::vector & Model::getMechanicalJointNames() const { return mechanicalJointNames_; @@ -1890,36 +1871,28 @@ namespace jiminy return backlashJointIndices_; } - /// \brief Returns true if at least one constraint is active on the robot. - bool Model::hasConstraints() const + const std::vector & Model::getLogPositionFieldnames() const { - bool hasConstraintsEnabled = false; - const_cast(constraints_) - .foreach( - [&hasConstraintsEnabled]( - const std::shared_ptr & constraint, - ConstraintNodeType /* node */) - { - if (constraint->getIsEnabled()) - { - hasConstraintsEnabled = true; - } - }); - return hasConstraintsEnabled; + return logPositionFieldnames_; } - Eigen::Index Model::nq() const + const std::vector & Model::getLogVelocityFieldnames() const { - return nq_; + return logVelocityFieldnames_; } - Eigen::Index Model::nv() const + const std::vector & Model::getLogAccelerationFieldnames() const { - return nv_; + return logAccelerationFieldnames_; } - Eigen::Index Model::nx() const + const std::vector & Model::getLogEffortFieldnames() const { - return nx_; + return logEffortFieldnames_; + } + + const std::vector & Model::getLogForceExternalFieldnames() const + { + return logForceExternalFieldnames_; } } diff --git a/core/src/robot/robot.cc b/core/src/robot/robot.cc index be1c624f6..941e07774 100644 --- a/core/src/robot/robot.cc +++ b/core/src/robot/robot.cc @@ -1,19 +1,22 @@ #include #include -#include "jiminy/core/utilities/helpers.h" -#include "jiminy/core/utilities/pinocchio.h" -#include "jiminy/core/utilities/json.h" -#include "jiminy/core/io/file_device.h" +#include "jiminy/core/constraints/abstract_constraint.h" +#include "jiminy/core/constraints/joint_constraint.h" #include "jiminy/core/hardware/abstract_motor.h" #include "jiminy/core/hardware/abstract_sensor.h" #include "jiminy/core/control/abstract_controller.h" #include "jiminy/core/control/controller_functor.h" -#include "jiminy/core/constraints/abstract_constraint.h" -#include "jiminy/core/constraints/joint_constraint.h" +#include "jiminy/core/io/file_device.h" +#include "jiminy/core/io/serialization.h" +#include "jiminy/core/utilities/helpers.h" +#include "jiminy/core/utilities/pinocchio.h" +#include "jiminy/core/utilities/json.h" #include "jiminy/core/robot/robot.h" +#include + namespace jiminy { @@ -26,6 +29,33 @@ namespace jiminy setOptions(getOptions()); } + Robot::Robot(const Robot & other) : + Model() + { + *this = other; + } + + Robot & Robot::operator=(const Robot & other) + { + // Serialize the user-specified robot + std::stringstream sstream; + { + stateful_binary_oarchive oa(sstream); + const Robot * const otherPtr = &other; + oa << boost::serialization::make_nvp("px", otherPtr); + } + + // De-serialize the user-specified robot then move assign to this robot + { + stateful_binary_iarchive ia(sstream); + Robot * robotPtr; + ia >> boost::serialization::make_nvp("px", robotPtr); + *this = std::move(*robotPtr); + } + + return *this; + } + Robot::~Robot() { // Detach all the motors and sensors @@ -201,6 +231,7 @@ namespace jiminy /* Trigger extended model regeneration. * The whole robot must be reset as joint and frame indices would be corrupted. + * Some sensors must be reset as their attributes may depends on motors options. * Skip reset if nothing has changed from this motor point of view. This is necessary to prevent infinite recursive reset loop. */ if (hasChanged) @@ -209,23 +240,20 @@ namespace jiminy } // Update rotor inertia - const std::string & mechanicaljointName = motorIn.getJointName(); - std::string inertiaJointName = mechanicaljointName; - const Eigen::Index motorVelocityIndex = motorIn.getJointVelocityIndex(); + const pinocchio::JointIndex jointIndex = motorIn.getJointIndex(); + const Eigen::Index motorVelocityIndex = + robot->pinocchioModel_.joints[jointIndex].idx_v(); robot->pinocchioModel_.rotorInertia[motorVelocityIndex] += motorIn.getArmature(); - - // Update effort limit - robot->pinocchioModel_.effortLimit[motorVelocityIndex] = motorIn.getCommandLimit(); }; - // Attach the motor motor->attach(shared_from_this(), notifyRobot, motorSharedStorage_.get()); - // Add the motor to the holder + // Add the motor to the registry motors_.push_back(motor); + motorsWeakConstRef_.push_back(motor); - // Refresh the motors proxies - refreshMotorProxies(); + // Append log telemetry motor command fieldnames + logCommandFieldnames_.push_back(toString(JOINT_PREFIX_BASE, "Command", motorName)); } void Robot::detachMotor(const std::string & motorName, bool triggerReset) @@ -257,9 +285,13 @@ namespace jiminy // Hold the motor alive before removing it std::shared_ptr motor = *motorIt; - // Remove the motor from the holder + // Remove the motor from registry + motorsWeakConstRef_.erase(motorsWeakConstRef_.begin() + motor->getIndex()); motors_.erase(motorIt); + // Remove log telemetry motor command fieldname + logCommandFieldnames_.erase(logCommandFieldnames_.begin() + motor->getIndex()); + // Detach the motor motor->detach(); @@ -271,16 +303,13 @@ namespace jiminy { (void)shared_from_this(); } - catch (std::bad_weak_ptr & e) + catch (std::bad_weak_ptr &) { return; } // Trigger extended model regeneration reset(std::random_device{}); - - // Refresh the motors proxies - refreshMotorProxies(); } } @@ -290,39 +319,53 @@ namespace jiminy detachMotor(motorName, true); } + template + static std::vector + getHardwareNames(const std::vector> & hardwares) + { + std::vector hardwareNames; + hardwareNames.reserve(hardwares.size()); + std::transform(hardwares.begin(), + hardwares.end(), + std::back_inserter(hardwareNames), + [](const std::shared_ptr & hardware) { return hardware->getName(); }); + return hardwareNames; + } + void Robot::detachMotors(std::vector motorNames) { + // Get the name of all motors currently attache to the robot + std::vector allMotorNames = getHardwareNames(motors_); + + // Remove all sensors if none is specified if (motorNames.empty()) { - // Remove all sensors if none is specified - if (!motorNames_.empty()) + if (!allMotorNames.empty()) { - detachMotors(motorNames_); + detachMotors(allMotorNames); } + return; } - else + + // Make sure that no motor names are duplicates + if (checkDuplicates(motorNames)) { - // Make sure that no motor names are duplicates - if (checkDuplicates(motorNames)) - { - JIMINY_THROW(std::invalid_argument, "Duplicated motor names found."); - } + JIMINY_THROW(std::invalid_argument, "Duplicated motor names found."); + } - // Make sure that every motor name exist - if (!checkInclusion(motorNames_, motorNames)) - { - JIMINY_THROW(std::invalid_argument, - "At least one of the motor names does not exist."); - } + // Make sure that every motor name exist + if (!checkInclusion(allMotorNames, motorNames)) + { + JIMINY_THROW(std::invalid_argument, "At least one of the motor names does not exist."); + } - /* Detach motors one-by-one, and reset proxies at the very end only. - Beware the motor names must be stored locally because the list - of motors will be updated in-place. */ - const std::size_t nMotors = motorNames.size(); - for (std::size_t i = 0; i < nMotors; ++i) - { - detachMotor(motorNames[i], i == (nMotors - 1)); - } + /* Detach motors one-by-one, and reset proxies at the very end only. + Beware the motor names must be stored locally because the list + of motors will be updated in-place. */ + const std::size_t nMotors = motorNames.size(); + for (std::size_t i = 0; i < nMotors; ++i) + { + detachMotor(motorNames[i], i == (nMotors - 1)); } } @@ -374,9 +417,7 @@ namespace jiminy // Create the sensor and add it to its group sensors_[sensorType].push_back(sensor); - - // Refresh the sensors proxies - refreshSensorProxies(); + sensorsWeakConstRef_[sensorType].push_back(sensor); } void Robot::detachSensor(const std::string & sensorType, const std::string & sensorName) @@ -420,24 +461,27 @@ namespace jiminy "'."); } - // Detach the sensor - (*sensorIt)->detach(); + // Hold the sensor alive before removing it + std::shared_ptr sensor = *sensorIt; // Remove the sensor from its group + WeakSensorVector & sensorGroupWeakConstRef = sensorsWeakConstRef_[sensorType]; + sensorGroupWeakConstRef.erase(sensorGroupWeakConstRef.begin() + sensor->getIndex()); sensorGroupIt->second.erase(sensorIt); + // Detach the sensor + sensor->detach(); + // Remove the sensor group if there is no more sensors left GenericConfig & telemetryOptions = boost::get(robotOptionsGeneric_.at("telemetry")); if (sensorGroupIt->second.empty()) { sensors_.erase(sensorType); + sensorsWeakConstRef_.erase(sensorType); sensorSharedStorageMap_.erase(sensorType); telemetryOptions.erase(sensorType); } - - // Refresh the sensors proxies - refreshSensorProxies(); } void Robot::detachSensors(const std::string & sensorType) @@ -453,8 +497,7 @@ namespace jiminy "'."); } - std::vector sensorGroupNames = - sensorNames_[sensorType]; // Make a copy since calling detachSensors update it ! + std::vector sensorGroupNames = getHardwareNames(sensorGroupIt->second); for (const std::string & sensorName : sensorGroupNames) { detachSensor(sensorType, sensorName); @@ -536,77 +579,6 @@ namespace jiminy return std::const_pointer_cast(controller_); } - void Robot::refreshProxies() - { - if (!isInitialized_) - { - JIMINY_THROW(bad_control_flow, "Robot not initialized."); - } - - Model::refreshProxies(); - refreshMotorProxies(); - refreshSensorProxies(); - } - - void Robot::refreshMotorProxies() - { - if (!isInitialized_) - { - JIMINY_THROW(bad_control_flow, "Robot not initialized."); - } - - // Determine the number of motors - nmotors_ = motors_.size(); - - // Extract the motor names - motorNames_.clear(); - motorNames_.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(motorNames_), - [](const auto & elem) -> std::string { return elem->getName(); }); - - // Generate the fieldnames associated with command - logCommandFieldnames_.clear(); - logCommandFieldnames_.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(logCommandFieldnames_), - [](const auto & elem) -> std::string - { return toString(JOINT_PREFIX_BASE, "Command", elem->getName()); }); - - // Generate the fieldnames associated with motor efforts - logMotorEffortFieldnames_.clear(); - logMotorEffortFieldnames_.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(logMotorEffortFieldnames_), - [](const auto & elem) -> std::string - { return toString(JOINT_PREFIX_BASE, "Effort", elem->getName()); }); - } - - void Robot::refreshSensorProxies() - { - if (!isInitialized_) - { - JIMINY_THROW(bad_control_flow, "Robot not initialized."); - } - - // Extract the motor names - sensorNames_.clear(); - sensorNames_.reserve(sensors_.size()); - for (const auto & [sensorType, sensorGroup] : sensors_) - { - std::vector sensorGroupNames; - sensorGroupNames.reserve(sensorGroup.size()); - std::transform(sensorGroup.begin(), - sensorGroup.end(), - std::back_inserter(sensorGroupNames), - [](const auto & elem) -> std::string { return elem->getName(); }); - sensorNames_.emplace(sensorType, std::move(sensorGroupNames)); - } - } - void Robot::initializeExtendedModel() { // Call base implementation @@ -695,11 +667,16 @@ namespace jiminy return std::const_pointer_cast(*motorIt); } - const Robot::MotorVector & Robot::getMotors() const + const Robot::MotorVector & Robot::getMotors() { return motors_; } + const Robot::WeakMotorVector & Robot::getMotors() const + { + return motorsWeakConstRef_; + } + std::shared_ptr Robot::getSensor(const std::string & sensorType, const std::string & sensorName) { @@ -768,11 +745,16 @@ namespace jiminy return std::const_pointer_cast(*sensorIt); } - const Robot::SensorTree & Robot::getSensors() const + const Robot::SensorTree & Robot::getSensors() { return sensors_; } + const Robot::WeakSensorTree & Robot::getSensors() const + { + return sensorsWeakConstRef_; + } + void Robot::setOptions(const GenericConfig & robotOptions) { // Make sure that no simulation is running @@ -934,19 +916,21 @@ namespace jiminy } } - const Eigen::VectorXd & Robot::getMotorEfforts() const + std::tuple, Eigen::Ref> + Robot::getMotorEfforts() const { static const Eigen::VectorXd motorsEffortsEmpty; + static const Eigen::VectorXd jointsEffortsEmpty; if (!motors_.empty()) { return (*motors_.begin())->getAll(); } - return motorsEffortsEmpty; + return {motorsEffortsEmpty, jointsEffortsEmpty}; } - double Robot::getMotorEffort(const std::string & motorName) const + std::tuple Robot::getMotorEffort(const std::string & motorName) const { if (!isInitialized_) { @@ -1059,6 +1043,11 @@ namespace jiminy controller_->updateTelemetry(); } + const std::vector & Robot::getLogCommandFieldnames() const + { + return logCommandFieldnames_; + } + std::unique_ptr Robot::getLock() { // Make sure that the robot is not already locked @@ -1081,94 +1070,8 @@ namespace jiminy return mutexLocal_->isLocked(); } - const std::vector & Robot::getMotorNames() const - { - return motorNames_; - } - - std::vector Robot::getMotorJointIndices() const - { - std::vector motorJointIndices; - motorJointIndices.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(motorJointIndices), - [](const auto & motor) -> pinocchio::JointIndex - { return motor->getJointIndex(); }); - return motorJointIndices; - } - - std::vector> Robot::getMotorsPositionIndices() const - { - std::vector> motorPositionIndices; - motorPositionIndices.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(motorPositionIndices), - [](const auto & elem) -> std::vector - { - const Eigen::Index & jointPositionIndex = elem->getJointPositionIndex(); - if (elem->getJointType() == JointModelType::ROTARY_UNBOUNDED) - { - return {jointPositionIndex, jointPositionIndex + 1}; - } - else - { - return {jointPositionIndex}; - } - }); - return motorPositionIndices; - } - - std::vector Robot::getMotorVelocityIndices() const - { - std::vector motorVelocityIndices; - motorVelocityIndices.reserve(nmotors_); - std::transform(motors_.begin(), - motors_.end(), - std::back_inserter(motorVelocityIndices), - [](const auto & elem) -> Eigen::Index - { return elem->getJointVelocityIndex(); }); - return motorVelocityIndices; - } - - const Eigen::VectorXd & Robot::getCommandLimit() const - { - return pinocchioModel_.effortLimit; - } - - const std::unordered_map> & Robot::getSensorNames() const - { - return sensorNames_; - } - - const std::vector & Robot::getSensorNames(const std::string & sensorType) const - { - static const std::vector sensorNamesEmpty{}; - - auto sensorsNamesIt = sensorNames_.find(sensorType); - if (sensorsNamesIt != sensorNames_.end()) - { - return sensorsNamesIt->second; - } - else - { - return sensorNamesEmpty; - } - } - - const std::vector & Robot::getLogCommandFieldnames() const - { - return logCommandFieldnames_; - } - - const std::vector & Robot::getLogMotorEffortFieldnames() const - { - return logMotorEffortFieldnames_; - } - Eigen::Index Robot::nmotors() const { - return nmotors_; + return motors_.size(); } } diff --git a/core/src/utilities/geometry.cc b/core/src/utilities/geometry.cc index f242d0b6f..8b11eca07 100644 --- a/core/src/utilities/geometry.cc +++ b/core/src/utilities/geometry.cc @@ -749,4 +749,50 @@ namespace jiminy } }; } + + HeightmapFunction stairs( + double stepWidth, double stepHeight, uint32_t stepNumber, double orientation) + { + const double interpDelta = 0.01; + const Eigen::Rotation2D rot_mat(orientation); + + return [stepWidth, stepHeight, stepNumber, rot_mat, interpDelta]( + const Eigen::Vector2d & pos, double & height, Eigen::Vector3d & normal) -> void + { + // Compute position in stairs reference frame + Eigen::Vector2d posRel = rot_mat.inverse() * pos; + const double modPos = std::fmod(std::abs(posRel[0]), stepWidth * stepNumber * 2); + + // Compute the default height and normal + uint32_t stairIndex = static_cast(modPos / stepWidth); + int8_t staircaseSlopeSign = 1; + if (stairIndex >= stepNumber) + { + stairIndex = 2 * stepNumber - stairIndex; + staircaseSlopeSign = -1; + } + height = stairIndex * stepHeight; + normal = Eigen::Vector3d::UnitZ(); + + // Avoid unsupported vertical edge + const double posRelOnStep = std::fmod(modPos, stepWidth) / stepWidth; + if (1.0 - posRelOnStep < interpDelta) + { + const double slope = staircaseSlopeSign * stepHeight / interpDelta; + // Update height + height += slope * (posRelOnStep - (1.0 - interpDelta)); + + // Compute the inverse of the normal's Euclidean norm + const double normInv = 1.0 / std::sqrt(1.0 + std::pow(slope, 2)); + + // Update normal vector + // step 1. compute normal in stairs reference frame: + // normal << -slope * normInv, 0.0, normInv; + // step 2. Rotate normal vector in world plane reference frame: + // normal.head<2>() = rot_mat * normal.head<2>(); + // Or simply in a single operation: + normal << -slope * normInv * rot_mat.toRotationMatrix().col(0), normInv; + } + }; + } } \ No newline at end of file diff --git a/core/src/utilities/pinocchio.cc b/core/src/utilities/pinocchio.cc index efc1af72a..e16b9da73 100644 --- a/core/src/utilities/pinocchio.cc +++ b/core/src/utilities/pinocchio.cc @@ -616,13 +616,22 @@ namespace jiminy pinocchio::FrameIndex childFrameIndex = i; do { - childFrameIndex = model.frames[childFrameIndex].previousFrame; + const pinocchio::FrameIndex previousFrameIndex = + model.frames[childFrameIndex].previousFrame; + if (childFrameIndex > 1 && previousFrameIndex == childFrameIndex) + { + JIMINY_THROW(std::runtime_error, + "There is something wrong with the model. Frame '", + model.frames[childFrameIndex].name, + "' is its own parent."); + } + childFrameIndex = previousFrameIndex; if (childFrameIndex == frameIndex) { childFrameIndices.push_back(i); break; } - } while (childFrameIndex > 0 && + } while (childFrameIndex > 1 && model.frames[childFrameIndex].type != pinocchio::FrameType::JOINT); } diff --git a/core/unit/engine_sanity_check.cc b/core/unit/engine_sanity_check.cc index ff0cab20c..429f46d75 100755 --- a/core/unit/engine_sanity_check.cc +++ b/core/unit/engine_sanity_check.cc @@ -64,17 +64,22 @@ TEST(EngineSanity, EnergyConservation) // Get all robot options GenericConfig robotOptions = robot->getOptions(); - // Disable velocity and position limits + // Disable position limits GenericConfig & modelOptions = boost::get(robotOptions.at("model")); GenericConfig & jointsOptions = boost::get(modelOptions.at("joints")); - boost::get(jointsOptions.at("enableVelocityLimit")) = false; + boost::get(jointsOptions.at("positionLimitFromUrdf")) = false; + boost::get(jointsOptions.at("positionLimitMin")) = + Eigen::Vector2d::Constant(-INF); + boost::get(jointsOptions.at("positionLimitMax")) = + Eigen::Vector2d::Constant(+INF); - // Disable torque limits + // Disable velocity and torque limits GenericConfig & motorsOptions = boost::get(robotOptions.at("motors")); for (auto & motorOptionsItem : motorsOptions) { GenericConfig & motorOptions = boost::get(motorOptionsItem.second); - boost::get(motorOptions.at("enableCommandLimit")) = false; + boost::get(motorOptions.at("enableVelocityLimit")) = false; + boost::get(motorOptions.at("enableEffortLimit")) = false; } // Set all robot options @@ -88,13 +93,17 @@ TEST(EngineSanity, EnergyConservation) Engine engine{}; engine.addRobot(robot); - // Configure engine: High accuracy + Continuous-time integration + // Configure engine: High accuracy + Continuous-time integration + telemetry GenericConfig simuOptions = engine.getDefaultEngineOptions(); { GenericConfig & stepperOptions = boost::get(simuOptions.at("stepper")); boost::get(stepperOptions.at("tolAbs")) = TOLERANCE * 1.0e-2; boost::get(stepperOptions.at("tolRel")) = TOLERANCE * 1.0e-2; } + { + GenericConfig & telemetryOptions = boost::get(simuOptions.at("telemetry")); + boost::get(telemetryOptions.at("enableEnergy")) = true; + } engine.setOptions(simuOptions); // Run simulation @@ -129,6 +138,10 @@ TEST(EngineSanity, EnergyConservation) boost::get(stepperOptions.at("sensorsUpdatePeriod")) = 1.0e-3; boost::get(stepperOptions.at("controllerUpdatePeriod")) = 1.0e-3; } + { + GenericConfig & telemetryOptions = boost::get(simuOptions.at("telemetry")); + boost::get(telemetryOptions.at("enableEnergy")) = true; + } engine.setOptions(simuOptions); // Run simulation diff --git a/core/unit/model_test.cc b/core/unit/model_test.cc index 02cf5d84a..07454d2b2 100644 --- a/core/unit/model_test.cc +++ b/core/unit/model_test.cc @@ -30,11 +30,12 @@ TEST_P(ModelTestFixture, CreateFlexible) model->initialize(urdfPath, hasFreeflyer, std::vector(), true); // We now have a rigid robot: perform rigid computation on this - auto q = pinocchio::randomConfiguration(model->pinocchioModel_); if (hasFreeflyer) { - q.head<3>().setZero(); + model->pinocchioModel_.lowerPositionLimit.head<3>().setZero(); + model->pinocchioModel_.upperPositionLimit.head<3>().setZero(); } + auto q = pinocchio::randomConfiguration(model->pinocchioModel_); auto pinocchioData = pinocchio::Data(model->pinocchioModel_); pinocchio::framesForwardKinematics(model->pinocchioModel_, pinocchioData, q); diff --git a/data/bipedal_robots/atlas/atlas_v4.urdf b/data/bipedal_robots/atlas/atlas.urdf similarity index 72% rename from data/bipedal_robots/atlas/atlas_v4.urdf rename to data/bipedal_robots/atlas/atlas.urdf index 702e7f656..333758bc2 100644 --- a/data/bipedal_robots/atlas/atlas_v4.urdf +++ b/data/bipedal_robots/atlas/atlas.urdf @@ -1,42 +1,40 @@ - - - - + + + + + + - - - + + + - + - + - - - - - - - - - - - - + + + + + + + + + + + + - - - - - + @@ -65,9 +63,9 @@ - - - + + + @@ -90,7 +88,7 @@ - + @@ -113,9 +111,9 @@ - - - + + + @@ -138,9 +136,9 @@ - - - + + + @@ -158,7 +156,7 @@ - + @@ -182,7 +180,7 @@ - + @@ -212,7 +210,7 @@ - + @@ -230,7 +228,7 @@ - + @@ -266,7 +264,7 @@ - + @@ -290,7 +288,7 @@ - + @@ -320,7 +318,7 @@ - + @@ -344,7 +342,7 @@ - + @@ -389,7 +387,7 @@ - + @@ -408,7 +406,7 @@ - + @@ -426,7 +424,7 @@ - + @@ -468,7 +466,7 @@ - + @@ -497,9 +495,9 @@ - - - + + + @@ -522,7 +520,7 @@ - + @@ -545,9 +543,9 @@ - - - + + + @@ -570,9 +568,9 @@ - - - + + + @@ -590,7 +588,7 @@ - + @@ -614,7 +612,7 @@ - + @@ -644,7 +642,7 @@ - + @@ -662,7 +660,7 @@ - + @@ -698,7 +696,7 @@ - + @@ -722,7 +720,7 @@ - + @@ -752,7 +750,7 @@ - + @@ -776,7 +774,7 @@ - + @@ -818,11 +816,10 @@ - - + @@ -928,7 +925,7 @@ - + @@ -937,7 +934,7 @@ - + @@ -946,7 +943,7 @@ - + @@ -955,25 +952,25 @@ - + - + - + - + - + @@ -989,7 +986,6 @@ - @@ -997,8 +993,8 @@ - - + + @@ -1006,40 +1002,38 @@ - - + + - + - - + + - + - - + - + - - + - + @@ -1048,7 +1042,7 @@ - + @@ -1057,7 +1051,7 @@ - + @@ -1066,43 +1060,43 @@ - + - + - + - + - + - + - + - + - + @@ -1111,9 +1105,8 @@ - + - @@ -1121,22 +1114,22 @@ - + - - + + - + - - + + @@ -1144,31 +1137,29 @@ - - + + - + - - + - + - - + - + @@ -1177,7 +1168,7 @@ - + @@ -1186,7 +1177,7 @@ - + @@ -1195,13 +1186,182 @@ - + - + + + + + 1000.0 + 1 + + /r_foot_contact + + + + + + + 1000.0 + 1 + + /l_foot_contact + + + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + + + + + 1 + - diff --git a/data/bipedal_robots/atlas/atlas_v4_hardware.toml b/data/bipedal_robots/atlas/atlas_hardware.toml similarity index 74% rename from data/bipedal_robots/atlas/atlas_v4_hardware.toml rename to data/bipedal_robots/atlas/atlas_hardware.toml index efe360431..8b1b16dff 100644 --- a/data/bipedal_robots/atlas/atlas_v4_hardware.toml +++ b/data/bipedal_robots/atlas/atlas_hardware.toml @@ -6,180 +6,240 @@ contactFrameNames = [] [Motor.SimpleMotor.neck_ry] joint_name = "neck_ry" +enableVelocityLimit = true +velocityEffortInvSlope = 0.08 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.back_bkz] joint_name = "back_bkz" +enableVelocityLimit = true +velocityEffortInvSlope = 0.03 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.back_bky] joint_name = "back_bky" +enableVelocityLimit = true +velocityEffortInvSlope = 0.005 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.back_bkx] joint_name = "back_bkx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.01 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_arm_shz] joint_name = "l_arm_shz" +enableVelocityLimit = true +velocityEffortInvSlope = 0.04 armature = 0.2 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_arm_shx] joint_name = "l_arm_shx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.04 armature = 0.2 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_arm_ely] joint_name = "l_arm_ely" +enableVelocityLimit = true +velocityEffortInvSlope = 0.06 armature = 0.2 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.l_arm_elx] joint_name = "l_arm_elx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.03 armature = 0.2 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.l_arm_wry] joint_name = "l_arm_wry" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_arm_wrx] joint_name = "l_arm_wrx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_arm_wry2] joint_name = "l_arm_wry2" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_arm_shz] joint_name = "r_arm_shz" +enableVelocityLimit = true +velocityEffortInvSlope = 0.04 armature = 0.2 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_arm_shx] joint_name = "r_arm_shx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.04 armature = 0.2 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_arm_ely] joint_name = "r_arm_ely" +enableVelocityLimit = true +velocityEffortInvSlope = 0.06 armature = 0.2 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.r_arm_elx] joint_name = "r_arm_elx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.03 armature = 0.2 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.r_arm_wry] joint_name = "r_arm_wry" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_arm_wrx] joint_name = "r_arm_wrx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_arm_wry2] joint_name = "r_arm_wry2" +enableVelocityLimit = true +velocityEffortInvSlope = 0.1 armature = 0.1 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_leg_hpx] joint_name = "l_leg_hpx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.005 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_leg_hpz] joint_name = "l_leg_hpz" +enableVelocityLimit = true +velocityEffortInvSlope = 0.01 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_leg_hpy] joint_name = "l_leg_hpy" +enableVelocityLimit = true +velocityEffortInvSlope = 0.004 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_leg_kny] joint_name = "l_leg_kny" +enableVelocityLimit = true +velocityEffortInvSlope = 0.004 armature = 0.6 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.l_leg_aky] joint_name = "l_leg_aky" +enableVelocityLimit = true +velocityEffortInvSlope = 0.005 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.l_leg_akx] joint_name = "l_leg_akx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.01 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_leg_hpx] joint_name = "r_leg_hpx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.005 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_leg_hpz] joint_name = "r_leg_hpz" +enableVelocityLimit = true +velocityEffortInvSlope = 0.01 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_leg_hpy] joint_name = "r_leg_hpy" +enableVelocityLimit = true +velocityEffortInvSlope = 0.004 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_leg_kny] joint_name = "r_leg_kny" +enableVelocityLimit = true +velocityEffortInvSlope = 0.004 armature = 0.6 frictionViscousPositive = -0.2 frictionViscousNegative = -0.2 [Motor.SimpleMotor.r_leg_aky] joint_name = "r_leg_aky" +enableVelocityLimit = true +velocityEffortInvSlope = 0.005 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 [Motor.SimpleMotor.r_leg_akx] joint_name = "r_leg_akx" +enableVelocityLimit = true +velocityEffortInvSlope = 0.01 armature = 0.6 frictionViscousPositive = -0.1 frictionViscousNegative = -0.1 @@ -197,186 +257,186 @@ frame_name = "l_foot" [Sensor.ForceSensor.r_foot] frame_name = "r_foot" -# ============= Effort sensors =============== +# ============= Encoder sensors =============== -[Sensor.EffortSensor.neck_ry] +[Sensor.EncoderSensor.neck_ry] motor_name = "neck_ry" -[Sensor.EffortSensor.back_bkz] +[Sensor.EncoderSensor.back_bkz] motor_name = "back_bkz" -[Sensor.EffortSensor.back_bky] +[Sensor.EncoderSensor.back_bky] motor_name = "back_bky" -[Sensor.EffortSensor.back_bkx] +[Sensor.EncoderSensor.back_bkx] motor_name = "back_bkx" -[Sensor.EffortSensor.l_arm_shz] +[Sensor.EncoderSensor.l_arm_shz] motor_name = "l_arm_shz" -[Sensor.EffortSensor.l_arm_shx] +[Sensor.EncoderSensor.l_arm_shx] motor_name = "l_arm_shx" -[Sensor.EffortSensor.l_arm_ely] +[Sensor.EncoderSensor.l_arm_ely] motor_name = "l_arm_ely" -[Sensor.EffortSensor.l_arm_elx] +[Sensor.EncoderSensor.l_arm_elx] motor_name = "l_arm_elx" -[Sensor.EffortSensor.l_arm_wry] +[Sensor.EncoderSensor.l_arm_wry] motor_name = "l_arm_wry" -[Sensor.EffortSensor.l_arm_wrx] +[Sensor.EncoderSensor.l_arm_wrx] motor_name = "l_arm_wrx" -[Sensor.EffortSensor.l_arm_wry2] +[Sensor.EncoderSensor.l_arm_wry2] motor_name = "l_arm_wry2" -[Sensor.EffortSensor.r_arm_shz] +[Sensor.EncoderSensor.r_arm_shz] motor_name = "r_arm_shz" -[Sensor.EffortSensor.r_arm_shx] +[Sensor.EncoderSensor.r_arm_shx] motor_name = "r_arm_shx" -[Sensor.EffortSensor.r_arm_ely] +[Sensor.EncoderSensor.r_arm_ely] motor_name = "r_arm_ely" -[Sensor.EffortSensor.r_arm_elx] +[Sensor.EncoderSensor.r_arm_elx] motor_name = "r_arm_elx" -[Sensor.EffortSensor.r_arm_wry] +[Sensor.EncoderSensor.r_arm_wry] motor_name = "r_arm_wry" -[Sensor.EffortSensor.r_arm_wrx] +[Sensor.EncoderSensor.r_arm_wrx] motor_name = "r_arm_wrx" -[Sensor.EffortSensor.r_arm_wry2] +[Sensor.EncoderSensor.r_arm_wry2] motor_name = "r_arm_wry2" -[Sensor.EffortSensor.l_leg_hpx] +[Sensor.EncoderSensor.l_leg_hpx] motor_name = "l_leg_hpx" -[Sensor.EffortSensor.l_leg_hpz] +[Sensor.EncoderSensor.l_leg_hpz] motor_name = "l_leg_hpz" -[Sensor.EffortSensor.l_leg_hpy] +[Sensor.EncoderSensor.l_leg_hpy] motor_name = "l_leg_hpy" -[Sensor.EffortSensor.l_leg_kny] +[Sensor.EncoderSensor.l_leg_kny] motor_name = "l_leg_kny" -[Sensor.EffortSensor.l_leg_aky] +[Sensor.EncoderSensor.l_leg_aky] motor_name = "l_leg_aky" -[Sensor.EffortSensor.l_leg_akx] +[Sensor.EncoderSensor.l_leg_akx] motor_name = "l_leg_akx" -[Sensor.EffortSensor.r_leg_hpx] +[Sensor.EncoderSensor.r_leg_hpx] motor_name = "r_leg_hpx" -[Sensor.EffortSensor.r_leg_hpz] +[Sensor.EncoderSensor.r_leg_hpz] motor_name = "r_leg_hpz" -[Sensor.EffortSensor.r_leg_hpy] +[Sensor.EncoderSensor.r_leg_hpy] motor_name = "r_leg_hpy" -[Sensor.EffortSensor.r_leg_kny] +[Sensor.EncoderSensor.r_leg_kny] motor_name = "r_leg_kny" -[Sensor.EffortSensor.r_leg_aky] +[Sensor.EncoderSensor.r_leg_aky] motor_name = "r_leg_aky" -[Sensor.EffortSensor.r_leg_akx] +[Sensor.EncoderSensor.r_leg_akx] motor_name = "r_leg_akx" -# ============= Encoder sensors =============== +# ============= Effort sensors =============== -[Sensor.EncoderSensor.neck_ry] -joint_name = "neck_ry" +[Sensor.EffortSensor.neck_ry] +motor_name = "neck_ry" -[Sensor.EncoderSensor.back_bkz] -joint_name = "back_bkz" +[Sensor.EffortSensor.back_bkz] +motor_name = "back_bkz" -[Sensor.EncoderSensor.back_bky] -joint_name = "back_bky" +[Sensor.EffortSensor.back_bky] +motor_name = "back_bky" -[Sensor.EncoderSensor.back_bkx] -joint_name = "back_bkx" +[Sensor.EffortSensor.back_bkx] +motor_name = "back_bkx" -[Sensor.EncoderSensor.l_arm_shz] -joint_name = "l_arm_shz" +[Sensor.EffortSensor.l_arm_shz] +motor_name = "l_arm_shz" -[Sensor.EncoderSensor.l_arm_shx] -joint_name = "l_arm_shx" +[Sensor.EffortSensor.l_arm_shx] +motor_name = "l_arm_shx" -[Sensor.EncoderSensor.l_arm_ely] -joint_name = "l_arm_ely" +[Sensor.EffortSensor.l_arm_ely] +motor_name = "l_arm_ely" -[Sensor.EncoderSensor.l_arm_elx] -joint_name = "l_arm_elx" +[Sensor.EffortSensor.l_arm_elx] +motor_name = "l_arm_elx" -[Sensor.EncoderSensor.l_arm_wry] -joint_name = "l_arm_wry" +[Sensor.EffortSensor.l_arm_wry] +motor_name = "l_arm_wry" -[Sensor.EncoderSensor.l_arm_wrx] -joint_name = "l_arm_wrx" +[Sensor.EffortSensor.l_arm_wrx] +motor_name = "l_arm_wrx" -[Sensor.EncoderSensor.l_arm_wry2] -joint_name = "l_arm_wry2" +[Sensor.EffortSensor.l_arm_wry2] +motor_name = "l_arm_wry2" -[Sensor.EncoderSensor.r_arm_shz] -joint_name = "r_arm_shz" +[Sensor.EffortSensor.r_arm_shz] +motor_name = "r_arm_shz" -[Sensor.EncoderSensor.r_arm_shx] -joint_name = "r_arm_shx" +[Sensor.EffortSensor.r_arm_shx] +motor_name = "r_arm_shx" -[Sensor.EncoderSensor.r_arm_ely] -joint_name = "r_arm_ely" +[Sensor.EffortSensor.r_arm_ely] +motor_name = "r_arm_ely" -[Sensor.EncoderSensor.r_arm_elx] -joint_name = "r_arm_elx" +[Sensor.EffortSensor.r_arm_elx] +motor_name = "r_arm_elx" -[Sensor.EncoderSensor.r_arm_wry] -joint_name = "r_arm_wry" +[Sensor.EffortSensor.r_arm_wry] +motor_name = "r_arm_wry" -[Sensor.EncoderSensor.r_arm_wrx] -joint_name = "r_arm_wrx" +[Sensor.EffortSensor.r_arm_wrx] +motor_name = "r_arm_wrx" -[Sensor.EncoderSensor.r_arm_wry2] -joint_name = "r_arm_wry2" +[Sensor.EffortSensor.r_arm_wry2] +motor_name = "r_arm_wry2" -[Sensor.EncoderSensor.l_leg_hpx] -joint_name = "l_leg_hpx" +[Sensor.EffortSensor.l_leg_hpx] +motor_name = "l_leg_hpx" -[Sensor.EncoderSensor.l_leg_hpz] -joint_name = "l_leg_hpz" +[Sensor.EffortSensor.l_leg_hpz] +motor_name = "l_leg_hpz" -[Sensor.EncoderSensor.l_leg_hpy] -joint_name = "l_leg_hpy" +[Sensor.EffortSensor.l_leg_hpy] +motor_name = "l_leg_hpy" -[Sensor.EncoderSensor.l_leg_kny] -joint_name = "l_leg_kny" +[Sensor.EffortSensor.l_leg_kny] +motor_name = "l_leg_kny" -[Sensor.EncoderSensor.l_leg_aky] -joint_name = "l_leg_aky" +[Sensor.EffortSensor.l_leg_aky] +motor_name = "l_leg_aky" -[Sensor.EncoderSensor.l_leg_akx] -joint_name = "l_leg_akx" +[Sensor.EffortSensor.l_leg_akx] +motor_name = "l_leg_akx" -[Sensor.EncoderSensor.r_leg_hpx] -joint_name = "r_leg_hpx" +[Sensor.EffortSensor.r_leg_hpx] +motor_name = "r_leg_hpx" -[Sensor.EncoderSensor.r_leg_hpz] -joint_name = "r_leg_hpz" +[Sensor.EffortSensor.r_leg_hpz] +motor_name = "r_leg_hpz" -[Sensor.EncoderSensor.r_leg_hpy] -joint_name = "r_leg_hpy" +[Sensor.EffortSensor.r_leg_hpy] +motor_name = "r_leg_hpy" -[Sensor.EncoderSensor.r_leg_kny] -joint_name = "r_leg_kny" +[Sensor.EffortSensor.r_leg_kny] +motor_name = "r_leg_kny" -[Sensor.EncoderSensor.r_leg_aky] -joint_name = "r_leg_aky" +[Sensor.EffortSensor.r_leg_aky] +motor_name = "r_leg_aky" -[Sensor.EncoderSensor.r_leg_akx] -joint_name = "r_leg_akx" \ No newline at end of file +[Sensor.EffortSensor.r_leg_akx] +motor_name = "r_leg_akx" diff --git a/data/bipedal_robots/atlas/atlas_v4_options.toml b/data/bipedal_robots/atlas/atlas_options.toml similarity index 83% rename from data/bipedal_robots/atlas/atlas_v4_options.toml rename to data/bipedal_robots/atlas/atlas_options.toml index 8a5bfa81a..26f21852b 100644 --- a/data/bipedal_robots/atlas/atlas_v4_options.toml +++ b/data/bipedal_robots/atlas/atlas_options.toml @@ -22,8 +22,3 @@ stabilizationFreq = 25.0 transitionEps = 5.0e-3 friction = 0.8 torsion = 0.0 - -# ======== Joints bounds configuration ======== - -[robot.model.joints] -enableVelocityLimit = true diff --git a/data/bipedal_robots/cassie/cassie_hardware.toml b/data/bipedal_robots/cassie/cassie_hardware.toml index 8e9f74407..42bf8b186 100644 --- a/data/bipedal_robots/cassie/cassie_hardware.toml +++ b/data/bipedal_robots/cassie/cassie_hardware.toml @@ -6,83 +6,103 @@ contactFrameNames = [] [Motor.SimpleMotor.hip_abduction_motor_left] joint_name = "hip_abduction_left" +enableVelocityLimit = true +velocityEffortInvSlope = 20.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 25.0 -armature = 0.2 # 0.038125 +armature = 0.00032 # 0.000061 [Motor.SimpleMotor.hip_rotation_motor_left] joint_name = "hip_rotation_left" +enableVelocityLimit = true +velocityEffortInvSlope = 20.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 25.0 -armature = 0.2 # 0.038125 +armature = 0.00032 # 0.000061 [Motor.SimpleMotor.hip_flexion_motor_left] joint_name = "hip_flexion_left" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.09344 +armature = 0.00078 # 0.000365 [Motor.SimpleMotor.knee_joint_motor_left] joint_name = "knee_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.09344 +armature = 0.00078 # 0.000365 [Motor.SimpleMotor.toe_joint_motor_left] joint_name = "toe_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.01225 +armature = 0.00008 # 0.0000049 [Motor.SimpleMotor.hip_abduction_motor_right] joint_name = "hip_abduction_right" +enableVelocityLimit = true +velocityEffortInvSlope = 20.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 25.0 -armature = 0.2 # 0.038125 +armature = 0.00032 # 0.000061 [Motor.SimpleMotor.hip_rotation_motor_right] joint_name = "hip_rotation_right" +enableVelocityLimit = true +velocityEffortInvSlope = 20.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 25.0 -armature = 0.2 # 0.038125 +armature = 0.00032 # 0.000061 [Motor.SimpleMotor.hip_flexion_motor_right] joint_name = "hip_flexion_right" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.09344 +armature = 0.00078 # 0.000365 [Motor.SimpleMotor.knee_joint_motor_right] joint_name = "knee_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.09344 +armature = 0.00078 # 0.000365 [Motor.SimpleMotor.toe_joint_motor_right] joint_name = "toe_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.01225 +armature = 0.00008 # 0.0000049 # =============== IMU sensors ================= @@ -132,31 +152,31 @@ motor_name = "toe_joint_motor_right" # ============= Encoder sensors =============== [Sensor.EncoderSensor.hip_abduction_left] -joint_name = "hip_abduction_left" +motor_name = "hip_abduction_motor_left" [Sensor.EncoderSensor.hip_rotation_left] -joint_name = "hip_rotation_left" +motor_name = "hip_rotation_motor_left" [Sensor.EncoderSensor.hip_flexion_left] -joint_name = "hip_flexion_left" +motor_name = "hip_flexion_motor_left" [Sensor.EncoderSensor.knee_joint_left] -joint_name = "knee_joint_left" +motor_name = "knee_joint_motor_left" [Sensor.EncoderSensor.toe_joint_left] -joint_name = "toe_joint_left" +motor_name = "toe_joint_motor_left" [Sensor.EncoderSensor.hip_abduction_right] -joint_name = "hip_abduction_right" +motor_name = "hip_abduction_motor_right" [Sensor.EncoderSensor.hip_rotation_right] -joint_name = "hip_rotation_right" +motor_name = "hip_rotation_motor_right" [Sensor.EncoderSensor.hip_flexion_right] -joint_name = "hip_flexion_right" +motor_name = "hip_flexion_motor_right" [Sensor.EncoderSensor.knee_joint_right] -joint_name = "knee_joint_right" +motor_name = "knee_joint_motor_right" [Sensor.EncoderSensor.toe_joint_right] -joint_name = "toe_joint_right" +motor_name = "toe_joint_motor_right" diff --git a/data/bipedal_robots/cassie/cassie_options.toml b/data/bipedal_robots/cassie/cassie_options.toml index 7ba101ac3..8e80e2bcb 100644 --- a/data/bipedal_robots/cassie/cassie_options.toml +++ b/data/bipedal_robots/cassie/cassie_options.toml @@ -22,7 +22,3 @@ stabilizationFreq = 25.0 transitionEps = 2.0e-3 friction = 0.5 -# ======== Joints bounds configuration ======== - -[robot.model.joints] -enableVelocityLimit = true diff --git a/data/bipedal_robots/digit/digit_hardware.toml b/data/bipedal_robots/digit/digit_hardware.toml index 4101a334c..73fc82691 100644 --- a/data/bipedal_robots/digit/digit_hardware.toml +++ b/data/bipedal_robots/digit/digit_hardware.toml @@ -6,6 +6,8 @@ contactFrameNames = [] [Motor.SimpleMotor.hip_abduction_left] joint_name = "hip_abduction_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -13,10 +15,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.hip_rotation_left] joint_name = "hip_rotation_left" +enableVelocityLimit = true +velocityEffortInvSlope = 130.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -24,10 +28,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.0675 +armature = 0.00008 # 0.000027 [Motor.SimpleMotor.hip_flexion_left] joint_name = "hip_flexion_left" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -35,10 +41,12 @@ frictionDryPositive = -0.5 frictionDryNegative = -0.5 frictionDrySlope = 20.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.120576 +armature = 0.00078125 # 0.00047 [Motor.SimpleMotor.knee_joint_left] joint_name = "knee_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -46,10 +54,12 @@ frictionDryPositive = -0.5 frictionDryNegative = -0.5 frictionDrySlope = 20.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.120576 +armature = 0.00078125 # 0.00047 [Motor.SimpleMotor.toe_pitch_joint_left] joint_name = "toe_pitch_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -57,10 +67,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.035 +armature = 0.00008 # 0.000014 [Motor.SimpleMotor.toe_roll_joint_left] joint_name = "toe_roll_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -68,10 +80,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.035 +armature = 0.00008 # 0.000014 [Motor.SimpleMotor.shoulder_roll_joint_left] joint_name = "shoulder_roll_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -2.0 frictionViscousNegative = -2.0 @@ -79,10 +93,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.shoulder_pitch_joint_left] joint_name = "shoulder_pitch_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -2.0 frictionViscousNegative = -2.0 @@ -90,10 +106,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.shoulder_yaw_joint_left] joint_name = "shoulder_yaw_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 100.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -101,10 +119,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.0675 +armature = 0.00008 # 0.000027 [Motor.SimpleMotor.elbow_joint_left] joint_name = "elbow_joint_left" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -112,10 +132,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.hip_abduction_right] joint_name = "hip_abduction_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -123,10 +145,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.hip_rotation_right] joint_name = "hip_rotation_right" +enableVelocityLimit = true +velocityEffortInvSlope = 130.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -134,10 +158,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.0675 +armature = 0.00008 # 0.000027 [Motor.SimpleMotor.hip_flexion_right] joint_name = "hip_flexion_right" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -145,10 +171,12 @@ frictionDryPositive = -0.5 frictionDryNegative = -0.5 frictionDrySlope = 20.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.120576 +armature = 0.00078125 # 0.00047 [Motor.SimpleMotor.knee_joint_right] joint_name = "knee_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 3.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -156,10 +184,12 @@ frictionDryPositive = -0.5 frictionDryNegative = -0.5 frictionDrySlope = 20.0 mechanicalReduction = 16.0 -armature = 0.2 # 0.120576 +armature = 0.00078125 # 0.00047 [Motor.SimpleMotor.toe_pitch_joint_right] joint_name = "toe_pitch_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -167,10 +197,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.035 +armature = 0.00008 # 0.000014 [Motor.SimpleMotor.toe_roll_joint_right] joint_name = "toe_roll_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -178,10 +210,12 @@ frictionDryPositive = -1.0 frictionDryNegative = -1.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.035 +armature = 0.00008 # 0.000014 [Motor.SimpleMotor.shoulder_roll_joint_right] joint_name = "shoulder_roll_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -2.0 frictionViscousNegative = -2.0 @@ -189,10 +223,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.shoulder_pitch_joint_right] joint_name = "shoulder_pitch_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 200.0 enableFriction = true frictionViscousPositive = -2.0 frictionViscousNegative = -2.0 @@ -200,10 +236,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 [Motor.SimpleMotor.shoulder_yaw_joint_right] joint_name = "shoulder_yaw_joint_right" +enableVelocityLimit = true +velocityEffortInvSlope = 100.0 enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -211,10 +249,12 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 50.0 -armature = 0.2 # 0.0675 +armature = 0.00008 # 0.000027 [Motor.SimpleMotor.elbow_joint_right] joint_name = "elbow_joint_right" +velocityEffortInvSlope = 200.0 +enableVelocityLimit = true enableFriction = true frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 @@ -222,7 +262,7 @@ frictionDryPositive = -2.0 frictionDryNegative = -2.0 frictionDrySlope = 20.0 mechanicalReduction = 80.0 -armature = 0.2 # 0.1728 +armature = 0.000031 # 0.000027 # =============== IMU sensors ================= @@ -240,64 +280,64 @@ frame_name = "right_toe_roll" # ============= Effort sensors =============== [Sensor.EncoderSensor.hip_abduction_left] -joint_name = "hip_abduction_left" +motor_name = "hip_abduction_left" [Sensor.EncoderSensor.hip_rotation_left] -joint_name = "hip_rotation_left" +motor_name = "hip_rotation_left" [Sensor.EncoderSensor.hip_flexion_left] -joint_name = "hip_flexion_left" +motor_name = "hip_flexion_left" [Sensor.EncoderSensor.knee_joint_left] -joint_name = "knee_joint_left" +motor_name = "knee_joint_left" [Sensor.EncoderSensor.toe_pitch_joint_left] -joint_name = "toe_pitch_joint_left" +motor_name = "toe_pitch_joint_left" [Sensor.EncoderSensor.toe_roll_joint_left] -joint_name = "toe_roll_joint_left" +motor_name = "toe_roll_joint_left" [Sensor.EncoderSensor.shoulder_roll_joint_left] -joint_name = "shoulder_roll_joint_left" +motor_name = "shoulder_roll_joint_left" [Sensor.EncoderSensor.shoulder_pitch_joint_left] -joint_name = "shoulder_pitch_joint_left" +motor_name = "shoulder_pitch_joint_left" [Sensor.EncoderSensor.shoulder_yaw_joint_left] -joint_name = "shoulder_yaw_joint_left" +motor_name = "shoulder_yaw_joint_left" [Sensor.EncoderSensor.elbow_joint_left] -joint_name = "elbow_joint_left" +motor_name = "elbow_joint_left" [Sensor.EncoderSensor.hip_abduction_right] -joint_name = "hip_abduction_right" +motor_name = "hip_abduction_right" [Sensor.EncoderSensor.hip_rotation_right] -joint_name = "hip_rotation_right" +motor_name = "hip_rotation_right" [Sensor.EncoderSensor.hip_flexion_right] -joint_name = "hip_flexion_right" +motor_name = "hip_flexion_right" [Sensor.EncoderSensor.knee_joint_right] -joint_name = "knee_joint_right" +motor_name = "knee_joint_right" [Sensor.EncoderSensor.toe_pitch_joint_right] -joint_name = "toe_pitch_joint_right" +motor_name = "toe_pitch_joint_right" [Sensor.EncoderSensor.toe_roll_joint_right] -joint_name = "toe_roll_joint_right" +motor_name = "toe_roll_joint_right" [Sensor.EncoderSensor.shoulder_roll_joint_right] -joint_name = "shoulder_roll_joint_right" +motor_name = "shoulder_roll_joint_right" [Sensor.EncoderSensor.shoulder_pitch_joint_right] -joint_name = "shoulder_pitch_joint_right" +motor_name = "shoulder_pitch_joint_right" [Sensor.EncoderSensor.shoulder_yaw_joint_right] -joint_name = "shoulder_yaw_joint_right" +motor_name = "shoulder_yaw_joint_right" [Sensor.EncoderSensor.elbow_joint_right] -joint_name = "elbow_joint_right" +motor_name = "elbow_joint_right" # ============= Effort sensors =============== diff --git a/data/bipedal_robots/digit/digit_options.toml b/data/bipedal_robots/digit/digit_options.toml index 7ba101ac3..5f8afc144 100644 --- a/data/bipedal_robots/digit/digit_options.toml +++ b/data/bipedal_robots/digit/digit_options.toml @@ -21,8 +21,3 @@ model = "constraint" stabilizationFreq = 25.0 transitionEps = 2.0e-3 friction = 0.5 - -# ======== Joints bounds configuration ======== - -[robot.model.joints] -enableVelocityLimit = true diff --git a/data/quadrupedal_robots/anymal/anymal_hardware.toml b/data/quadrupedal_robots/anymal/anymal_hardware.toml index b745d21e3..fe75c33d2 100644 --- a/data/quadrupedal_robots/anymal/anymal_hardware.toml +++ b/data/quadrupedal_robots/anymal/anymal_hardware.toml @@ -2,54 +2,87 @@ collisionBodyNames = [] contactFrameNames = ["LF_FOOT", "LH_FOOT", "RF_FOOT", "RH_FOOT"] +# ================== Motors ================== + [Motor.SimpleMotor.LF_HAA] joint_name = "LF_HAA" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.LF_HFE] joint_name = "LF_HFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.LF_KFE] joint_name = "LF_KFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RF_HAA] joint_name = "RF_HAA" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RF_HFE] joint_name = "RF_HFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RF_KFE] joint_name = "RF_KFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.LH_HAA] joint_name = "LH_HAA" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.LH_HFE] joint_name = "LH_HFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.LH_KFE] joint_name = "LH_KFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RH_HAA] joint_name = "RH_HAA" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RH_HFE] joint_name = "RH_HFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 [Motor.SimpleMotor.RH_KFE] joint_name = "RH_KFE" +enableVelocityLimit = true +velocityEffortInvSlope = 0.02 armature = 0.1 +# =============== IMU sensors ================= + +[Sensor.ImuSensor.imu_link] +frame_name = "imu_link" + +# ============= Force sensors =============== + [Sensor.ForceSensor.LF_FOOT] frame_name = "LF_FOOT" @@ -62,44 +95,45 @@ frame_name = "RF_FOOT" [Sensor.ForceSensor.RH_FOOT] frame_name = "RH_FOOT" -[Sensor.ImuSensor.imu_link] -frame_name = "imu_link" +# ============= Encoder sensors =============== [Sensor.EncoderSensor.LF_HAA] -joint_name = "LF_HAA" +motor_name = "LF_HAA" [Sensor.EncoderSensor.LF_HFE] -joint_name = "LF_HFE" +motor_name = "LF_HFE" [Sensor.EncoderSensor.LF_KFE] -joint_name = "LF_KFE" +motor_name = "LF_KFE" [Sensor.EncoderSensor.RF_HAA] -joint_name = "RF_HAA" +motor_name = "RF_HAA" [Sensor.EncoderSensor.RF_HFE] -joint_name = "RF_HFE" +motor_name = "RF_HFE" [Sensor.EncoderSensor.RF_KFE] -joint_name = "RF_KFE" +motor_name = "RF_KFE" [Sensor.EncoderSensor.LH_HAA] -joint_name = "LH_HAA" +motor_name = "LH_HAA" [Sensor.EncoderSensor.LH_HFE] -joint_name = "LH_HFE" +motor_name = "LH_HFE" [Sensor.EncoderSensor.LH_KFE] -joint_name = "LH_KFE" +motor_name = "LH_KFE" [Sensor.EncoderSensor.RH_HAA] -joint_name = "RH_HAA" +motor_name = "RH_HAA" [Sensor.EncoderSensor.RH_HFE] -joint_name = "RH_HFE" +motor_name = "RH_HFE" [Sensor.EncoderSensor.RH_KFE] -joint_name = "RH_KFE" +motor_name = "RH_KFE" + +# ============= Effort sensors =============== [Sensor.EffortSensor.LF_HAA] motor_name = "LF_HAA" diff --git a/data/quadrupedal_robots/anymal/anymal_options.toml b/data/quadrupedal_robots/anymal/anymal_options.toml index eeeeb6fc0..e43e89322 100644 --- a/data/quadrupedal_robots/anymal/anymal_options.toml +++ b/data/quadrupedal_robots/anymal/anymal_options.toml @@ -21,8 +21,3 @@ model = "constraint" stabilizationFreq = 25.0 transitionEps = 5.0e-3 friction = 1.0 - -# ======== Joints bounds configuration ======== - -[robot.model.joints] -enableVelocityLimit = true diff --git a/data/toys_models/ant/ant_hardware.toml b/data/toys_models/ant/ant_hardware.toml index 6141d2115..a7785e5cd 100644 --- a/data/toys_models/ant/ant_hardware.toml +++ b/data/toys_models/ant/ant_hardware.toml @@ -9,82 +9,89 @@ contactFrameNames = [] #['tip_1', 'tip_2', 'tip_3', 'tip_4'] [Motor.SimpleMotor.hip_1] joint_name = "hip_1" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.ankle_1] joint_name = "ankle_1" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.hip_2] joint_name = "hip_2" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.ankle_2] joint_name = "ankle_2" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.hip_3] joint_name = "hip_3" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.ankle_3] joint_name = "ankle_3" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.hip_4] joint_name = "hip_4" +enableVelocityLimit = false mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 [Motor.SimpleMotor.ankle_4] joint_name = "ankle_4" mechanicalReduction = 150.0 -armature = 1.0 +armature = 4.4e-05 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 # ============= Encoder sensors =============== [Sensor.EncoderSensor.hip_1] -joint_name = "hip_1" +motor_name = "hip_1" [Sensor.EncoderSensor.ankle_1] -joint_name = "ankle_1" +motor_name = "ankle_1" [Sensor.EncoderSensor.hip_2] -joint_name = "hip_2" +motor_name = "hip_2" [Sensor.EncoderSensor.ankle_2] -joint_name = "ankle_2" +motor_name = "ankle_2" [Sensor.EncoderSensor.hip_3] -joint_name = "hip_3" +motor_name = "hip_3" [Sensor.EncoderSensor.ankle_3] -joint_name = "ankle_3" +motor_name = "ankle_3" [Sensor.EncoderSensor.hip_4] -joint_name = "hip_4" +motor_name = "hip_4" [Sensor.EncoderSensor.ankle_4] -joint_name = "ankle_4" +motor_name = "ankle_4" diff --git a/data/toys_models/ant/ant_options.toml b/data/toys_models/ant/ant_options.toml index 4ae169e62..81b55cc46 100644 --- a/data/toys_models/ant/ant_options.toml +++ b/data/toys_models/ant/ant_options.toml @@ -19,8 +19,3 @@ model = "constraint" stabilizationFreq = 5.0 transitionEps = 1.0e-2 friction = 1.0 - -# ======== Joints bounds configuration ======== - -[robot.model.joints] -enableVelocityLimit = false diff --git a/data/toys_models/cartpole/cartpole.urdf b/data/toys_models/cartpole/cartpole.urdf index 1f60ec4ab..9cf636715 100644 --- a/data/toys_models/cartpole/cartpole.urdf +++ b/data/toys_models/cartpole/cartpole.urdf @@ -39,12 +39,12 @@ - + - + diff --git a/docs/api/gym_jiminy/common/bases/compositions.rst b/docs/api/gym_jiminy/common/bases/compositions.rst new file mode 100644 index 000000000..a3e8e45eb --- /dev/null +++ b/docs/api/gym_jiminy/common/bases/compositions.rst @@ -0,0 +1,6 @@ +Compositions +============ + +.. automodule:: gym_jiminy.common.bases.compositions + :members: + :show-inheritance: diff --git a/docs/api/gym_jiminy/common/bases/index.rst b/docs/api/gym_jiminy/common/bases/index.rst index a57f4f198..2147d8c96 100644 --- a/docs/api/gym_jiminy/common/bases/index.rst +++ b/docs/api/gym_jiminy/common/bases/index.rst @@ -6,4 +6,6 @@ Bases interfaces blocks + quantities + compositions pipeline diff --git a/docs/api/gym_jiminy/common/bases/quantities.rst b/docs/api/gym_jiminy/common/bases/quantities.rst new file mode 100644 index 000000000..f9e06101e --- /dev/null +++ b/docs/api/gym_jiminy/common/bases/quantities.rst @@ -0,0 +1,6 @@ +Quantities +========== + +.. automodule:: gym_jiminy.common.bases.quantities + :members: + :show-inheritance: diff --git a/docs/api/gym_jiminy/common/compositions/generic.rst b/docs/api/gym_jiminy/common/compositions/generic.rst new file mode 100644 index 000000000..ac3d2f006 --- /dev/null +++ b/docs/api/gym_jiminy/common/compositions/generic.rst @@ -0,0 +1,9 @@ +Generic +======= + +.. automodule:: gym_jiminy.common.compositions.generic + :members: + :undoc-members: + :private-members: + :inherited-members: + :show-inheritance: diff --git a/docs/api/gym_jiminy/common/compositions/index.rst b/docs/api/gym_jiminy/common/compositions/index.rst new file mode 100644 index 000000000..69854c700 --- /dev/null +++ b/docs/api/gym_jiminy/common/compositions/index.rst @@ -0,0 +1,8 @@ +Compositions +============ + +.. toctree:: + :maxdepth: 1 + + generic + locomotion diff --git a/docs/api/gym_jiminy/common/compositions/locomotion.rst b/docs/api/gym_jiminy/common/compositions/locomotion.rst new file mode 100644 index 000000000..17658944e --- /dev/null +++ b/docs/api/gym_jiminy/common/compositions/locomotion.rst @@ -0,0 +1,8 @@ +Locomotion +========== + +.. automodule:: gym_jiminy.common.compositions.locomotion + :members: + :undoc-members: + :private-members: + :show-inheritance: diff --git a/docs/api/gym_jiminy/common/index.rst b/docs/api/gym_jiminy/common/index.rst index 948096474..a71e72016 100644 --- a/docs/api/gym_jiminy/common/index.rst +++ b/docs/api/gym_jiminy/common/index.rst @@ -7,5 +7,7 @@ Gym Jiminy API bases/index blocks/index envs/index + quantities/index + rewards/index wrappers/index utils/index diff --git a/docs/api/gym_jiminy/common/quantities/generic.rst b/docs/api/gym_jiminy/common/quantities/generic.rst new file mode 100644 index 000000000..a2dafbbf3 --- /dev/null +++ b/docs/api/gym_jiminy/common/quantities/generic.rst @@ -0,0 +1,9 @@ +Generic +======= + +.. automodule:: gym_jiminy.common.quantities.generic + :members: + :undoc-members: + :private-members: + :inherited-members: + :show-inheritance: diff --git a/docs/api/gym_jiminy/common/quantities/index.rst b/docs/api/gym_jiminy/common/quantities/index.rst new file mode 100644 index 000000000..5fd849949 --- /dev/null +++ b/docs/api/gym_jiminy/common/quantities/index.rst @@ -0,0 +1,8 @@ +Quantities +========== + +.. toctree:: + :maxdepth: 1 + + generic + locomotion diff --git a/docs/api/gym_jiminy/common/quantities/locomotion.rst b/docs/api/gym_jiminy/common/quantities/locomotion.rst new file mode 100644 index 000000000..8e91bafc7 --- /dev/null +++ b/docs/api/gym_jiminy/common/quantities/locomotion.rst @@ -0,0 +1,8 @@ +Locomotion +========== + +.. automodule:: gym_jiminy.common.quantities.locomotion + :members: + :undoc-members: + :private-members: + :show-inheritance: diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/bases/__init__.py index bd5b503f8..2925c1d14 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/__init__.py @@ -1,8 +1,5 @@ # pylint: disable=missing-module-docstring -from .quantity import (SharedCache, - QuantityCreator, - AbstractQuantity) from .interfaces import (DT_EPS, ObsT, ActT, @@ -14,6 +11,16 @@ InterfaceObserver, InterfaceController, InterfaceJiminyEnv) +from .quantities import (QuantityCreator, + QuantityEvalMode, + SharedCache, + InterfaceQuantity, + AbstractQuantity, + StateQuantity, + DatasetTrajectoryQuantity) +from .compositions import (AbstractReward, + BaseQuantityReward, + BaseMixtureReward) from .blocks import (BlockStateT, InterfaceBlock, BaseObserverBlock, @@ -22,14 +29,12 @@ BasePipelineWrapper, BaseTransformObservation, BaseTransformAction, + ComposedJiminyEnv, ObservedJiminyEnv, ControlledJiminyEnv) __all__ = [ - 'SharedCache', - 'QuantityCreator', - 'AbstractQuantity', 'DT_EPS', 'ObsT', 'NestedObsT', @@ -40,15 +45,26 @@ 'InfoType', 'SensorMeasurementStackMap', 'EngineObsType', + 'SharedCache', 'InterfaceObserver', 'InterfaceController', 'InterfaceJiminyEnv', 'InterfaceBlock', + 'InterfaceQuantity', + 'AbstractQuantity', + 'AbstractReward', + 'BaseQuantityReward', + 'BaseMixtureReward', 'BaseObserverBlock', 'BaseControllerBlock', 'BasePipelineWrapper', 'BaseTransformObservation', 'BaseTransformAction', + 'ComposedJiminyEnv', 'ObservedJiminyEnv', 'ControlledJiminyEnv', + 'QuantityEvalMode', + 'QuantityCreator', + 'StateQuantity', + 'DatasetTrajectoryQuantity' ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/compositions.py b/python/gym_jiminy/common/gym_jiminy/common/bases/compositions.py new file mode 100644 index 000000000..c638aff03 --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/compositions.py @@ -0,0 +1,361 @@ +"""This module promotes reward components as first-class objects. + +Defining rewards this way allows for standardization of usual metrics. Overall, +it greatly reduces code duplication and bugs. +""" +from abc import ABC, abstractmethod +from typing import Sequence, Callable, Optional, Tuple, TypeVar + +import numpy as np + +from .interfaces import InfoType, InterfaceJiminyEnv +from .quantities import QuantityCreator + + +ValueT = TypeVar('ValueT') + + +class AbstractReward(ABC): + """Abstract class from which all reward component must derived. + + This goal of the agent is to maximize the expectation of the cumulative sum + of discounted reward over complete episodes. This holds true no matter if + its sign is always negative (aka. reward), always positive (aka. cost) or + indefinite (aka. objective). + + Defining cost is allowed by not recommended. Although it encourages the + agent to achieve the task at hands as quickly as possible if success if the + only termination condition, it has the side-effect to give the opportunity + to the agent to maximize the return by killing itself whenever this is an + option, which is rarely the desired behavior. No restriction is enforced as + it may be limiting in some relevant cases, so it is up to the user to make + sure that its design makes sense overall. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + name: str) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param name: Desired name of the reward. + """ + self.env = env + self._name = name + + @property + def name(self) -> str: + """Name uniquely identifying a given reward component. + + This name will be used as key for storing reward-specific monitoring + information in 'info' if key is missing, otherwise it will raise an + exception. + """ + return self._name + + @property + def is_terminal(self) -> Optional[bool]: + """Whether the reward is terminal, non-terminal, or indefinite. + + A reward is said to be "terminal" if only evaluated for the terminal + state of the MDP, "non-terminal" if evaluated for all states except the + terminal one, or indefinite if systematically evaluated no matter what. + + All rewards are supposed to be indefinite unless stated otherwise by + overloading this method. The responsibility of evaluating the reward + only when necessary is delegated to `compute`. This allows for complex + evaluation logics beyond terminal or non-terminal without restriction. + + .. note:: + Truncation is not consider the same as termination. The reward to + not be evaluated in such a case, which means that it will never be + for such episodes. + """ + return None + + @property + @abstractmethod + def is_normalized(self) -> bool: + """Whether the reward is guaranteed to be normalized, ie it is in range + [0.0, 1.0]. + """ + + @abstractmethod + def compute(self, terminated: bool, info: InfoType) -> Optional[float]: + """Compute the reward. + + .. note:: + Return value can be set to `None` to indicate that evaluation was + skipped for some reason, and therefore the reward must not be taken + into account when computing the total reward. This is useful when + the reward is undefined or simply inappropriate in the current + state of the environment. + + .. warning:: + It is the responsibility of the practitioner overloading this + method to honor flags 'is_terminated' (if not indefinite) and + 'is_normalized'. Failing this, an exception will be raised. + + :returns: Scalar value if the reward was evaluated, `None` otherwise. + """ + + def __call__(self, terminated: bool, info: InfoType) -> float: + """Return the reward associated with the current environment step. + + For the corresponding MDP to be stationary, the computation of the + reward is supposed to involve only the transition from previous to + current state of the environment under the ongoing action. + + .. note:: + This method is a lightweight wrapper around `compute` to skip + evaluation depending on whether the current state and the reward + are terminal. If the reward was truly evaluated, then 'info' is + updated to store either reward-specific 'info' if any or its value + otherwise. If not, then 'info' is left as-is and 0.0 is returned. + + .. warning:: + This method is not meant to be overloaded. + + :param terminated: Whether the episode has reached a terminal state of + the MDP at the current step. + :param info: Dictionary of extra information for monitoring. It will be + updated in-place for storing current value of the reward + in 'info' if it was truly evaluated. + """ + # Evaluate the reward and store extra information + reward_info: InfoType = {} + value = self.compute(terminated, reward_info) + + # Early return if None, which means that the reward was not evaluated + if value is None: + return 0.0 + + # Make sure that terminal flag is honored + if self.is_terminal is not None and self.is_terminal ^ terminated: + raise ValueError("Flag 'is_terminal' not honored.") + + # Make sure that the reward is scalar + assert np.ndim(value) == 0 + + # Make sure that the reward is normalized + if self.is_normalized and (value < 0.0 or value > 1.0): + raise ValueError( + "Reward not normalized in range [0.0, 1.0] as it ought to be.") + + # Store its value as info + if self.name is info.keys(): + raise KeyError( + f"Key '{self.name}' already reserved in 'info'. Impossible to " + "store value of reward component.") + if reward_info: + info[self.name] = reward_info + else: + info[self.name] = value + + # Returning the reward + return value + + +class BaseQuantityReward(AbstractReward): + """Base class that makes easy easy to derive reward components from generic + quantities. + + All this class does is applying some user-specified post-processing to the + value of a given multi-variate quantity to return a floating-point scalar + value, eventually normalized between 0.0 and 1.0 if desired. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + name: str, + quantity: QuantityCreator[ValueT], + transform_fn: Optional[Callable[[ValueT], float]], + is_normalized: bool, + is_terminal: Optional[bool]) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param name: Desired name of the reward. This name will be used as key + for storing current value of the reward in 'info', and to + add the underlying quantity to the set of already managed + quantities by the environment. As a result, it must be + unique otherwise an exception will be raised. + :param quantity: Tuple gathering the class of the underlying quantity + to use as reward after some post-processing, plus all + its constructor keyword-arguments except environment + 'env' and parent 'parent. + :param transform_fn: Transform function responsible for aggregating a + multi-variate quantity as floating-point scalar + value to maximize. Typical examples are `np.min`, + `np.max`, `lambda x: np.linalg.norm(x, order=N)`. + This function is also responsible for rescaling + the transformed quantity in range [0.0, 1.0] if + the reward is advertised as normalized. The Radial + Basis Function (RBF) kernel is the most common + choice to derive a reward to maximize from errors + based on distance metrics (See + `radial_basis_function` for details.). `None` to + skip transform entirely if not necessary. + :param is_normalized: Whether the reward is guaranteed to be normalized + after applying transform function `transform_fn`. + :param is_terminal: Whether the reward is terminal, non-terminal or + indefinite. A terminal reward will be evaluated at + most once, at the end of each episode for which a + termination condition has been triggered. On the + contrary, a non-terminal reward will be evaluated + systematically except at the end of the episode. + Finally, a indefinite reward will be evaluated + systematically. The value 0.0 is returned and no + 'info' will be stored when reward evaluation is + skipped. + """ + # Backup user argument(s) + self._transform_fn = transform_fn + self._is_normalized = is_normalized + self._is_terminal = is_terminal + + # Call base implementation + super().__init__(env, name) + + # Add quantity to the set of quantities managed by the environment + self.env.quantities[self.name] = quantity + + # Keep track of the underlying quantity + self.quantity = self.env.quantities.registry[self.name] + + @property + def is_terminal(self) -> Optional[bool]: + return self._is_terminal + + @property + def is_normalized(self) -> bool: + return self._is_normalized + + def compute(self, terminated: bool, info: InfoType) -> Optional[float]: + """Compute the reward if necessary depending on whether the reward and + state are terminal. If so, then first evaluate the underlying quantity, + next apply post-processing if requested. + + .. warning:: + This method is not meant to be overloaded. + + :returns: Scalar value if the reward was evaluated, `None` otherwise. + """ + # Early return depending on whether the reward and state are terminal + if self.is_terminal is not None and self.is_terminal ^ terminated: + return None + + # Evaluate raw quantity + value = self.env.quantities[self.name] + + # Early return if quantity is None + if value is None: + return None + + # Apply some post-processing if requested + if self._transform_fn is not None: + value = self._transform_fn(value) + + # Return the reward + return value + + +BaseQuantityReward.name.__doc__ = \ + """Name uniquely identifying every reward. + + It will be used as key not only for storing reward-specific monitoring + information in 'info', but also for adding the underlying quantity to + the ones already managed by the environment. + """ + + +class BaseMixtureReward(AbstractReward): + """Base class for aggregating multiple independent reward components in a + single one. + """ + + components: Tuple[AbstractReward, ...] + """List of all the reward components that must be aggregated together. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + name: str, + components: Sequence[AbstractReward], + reduce_fn: Callable[ + [Sequence[Optional[float]]], Optional[float]], + is_normalized: bool) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param name: Desired name of the total reward. + :param components: Sequence of reward components to aggregate. + :param reduce_fn: Transform function responsible for aggregating all + the reward components that were evaluated. Typical + examples are cumulative product and weighted sum. + :param is_normalized: Whether the reward is guaranteed to be normalized + after applying reduction function `reduce_fn`. + """ + # Make sure that at least one reward component has been specified + if not components: + raise ValueError( + "At least one reward component must be specified.") + + # Make sure that all reward components share the same environment + for reward in components: + if env is not reward.env: + raise ValueError( + "All reward components must share the same environment.") + + # Backup some user argument(s) + self.components = tuple(components) + self._reduce_fn = reduce_fn + self._is_normalized = is_normalized + + # Call base implementation + super().__init__(env, name) + + # Determine whether the reward mixture is terminal + is_terminal = {reward.is_terminal for reward in self.components} + self._is_terminal: Optional[bool] = None + if len(is_terminal) == 1: + self._is_terminal = next(iter(is_terminal)) + + @property + def is_terminal(self) -> Optional[bool]: + """Whether the reward is terminal, ie only evaluated at the end of an + episode if a termination condition has been triggered. + + The cumulative reward is considered terminal if and only if all its + individual reward components are terminal. + """ + return self._is_terminal + + @property + def is_normalized(self) -> bool: + return self._is_normalized + + def compute(self, terminated: bool, info: InfoType) -> Optional[float]: + """Evaluate each individual reward component for the current state of + the environment, then aggregate them in one. + """ + # Early return depending on whether the reward and state are terminal + if self.is_terminal is not None and self.is_terminal ^ terminated: + return None + + # Compute all reward components + values = [] + for reward in self.components: + # Evaluate reward + reward_info: InfoType = {} + value: Optional[float] = reward(terminated, reward_info) + + # Clear reward value if the reward was never truly evaluated + if not reward_info: + value = None + + # Append reward value and information + info.update(reward_info) + values.append(value) + + # Aggregate all reward components in one + reward_total = self._reduce_fn(values) + + return reward_total diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py index ebcd8bf91..f42efc7f7 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py @@ -137,31 +137,24 @@ def compute_command(self, action: ActT, command: BaseActT) -> None: def compute_reward(self, terminated: bool, # pylint: disable=unused-argument - truncated: bool, # pylint: disable=unused-argument info: InfoType # pylint: disable=unused-argument ) -> float: - """Compute the reward related to a specific control block. + """Compute the reward related to a specific control block, plus extra + information that may be helpful for monitoring or debugging purposes. For the corresponding MDP to be stationary, the computation of the reward is supposed to involve only the transition from previous to current state of the simulation (possibly comprising multiple agents) under the ongoing action. - By default, it returns 0.0 no matter what. It is up to the user to - provide a dedicated reward function whenever appropriate. - - .. warning:: - Only returning an aggregated scalar reward is supported. Yet, it is - possible to update 'info' by reference if one wants for keeping - track of individual reward components or any kind of extra info - that may be helpful for monitoring or debugging purposes. + By default, it returns 0.0 without extra information no matter what. + The user is expected to provide an appropriate reward on its own, + either by overloading this method or by wrapping the environment with + `ComposedJiminyEnv` for modular environment pipeline design. :param terminated: Whether the episode has reached the terminal state of the MDP at the current step. This flag can be used to compute a specific terminal reward. - :param truncated: Whether a truncation condition outside the scope of - the MDP has been satisfied at the current step. This - flag can be used to adapt the reward. :param info: Dictionary of extra information for monitoring. :returns: Aggregated reward for the current step. @@ -182,6 +175,7 @@ class InterfaceJiminyEnv( """Observer plus controller interface for both generic pipeline blocks, including environments. """ + metadata: Dict[str, Any] = { "render_modes": ( ['rgb_array'] + (['human'] if is_display_available() else [])) @@ -194,6 +188,16 @@ class InterfaceJiminyEnv( sensor_measurements: SensorMeasurementStackMap is_simulation_running: npt.NDArray[np.bool_] + num_steps: npt.NDArray[np.int64] + """Number of simulation steps that has been performed since last reset of + the base environment. + + .. note:: + The counter is incremented before updating the observation at the end + of the step, and consequently, before evaluating the reward and the + termination conditions. + """ + quantities: "QuantityManager" action: ActT @@ -250,13 +254,30 @@ def _observer_handle(self, :param v: Current extended velocity vector of the robot. :param sensor_measurements: Current sensor data. """ + # Early return if no simulation is running + if not self.is_simulation_running: + return + + # Reset the quantity manager. + # In principle, the internal cache of quantities should be cleared each + # time the state of the robot and/or its derivative changes. This is + # hard to do because there is no way to detect this specifically at the + # time being. However, `_observer_handle` is never called twice in the + # exact same state by the engine, so resetting quantities at the + # beginning of the method should cover most cases. Yet, quantities + # cannot be used reliably in the definition of profile forces because + # they are always updated before the controller gets called, no matter + # if either one or the other is time-continuous. Hacking the internal + # dynamics to clear quantities does not address this issue either. + self.quantities.clear() + # Refresh the observation if not already done but only if a simulation # is already running. It would be pointless to refresh the observation # at this point since the controller will be called multiple times at # start. Besides, it would defeat the purpose `_initialize_buffers`, # that is supposed to be executed before `refresh_observation` is being # called for the first time of an episode. - if not self.__is_observation_refreshed and self.is_simulation_running: + if not self.__is_observation_refreshed: measurement = self.__measurement measurement["t"][()] = t measurement["states"]["agent"]["q"] = q @@ -303,19 +324,6 @@ def _controller_handle(self, :returns: Motors torques to apply on the robot. """ - # Reset the quantity manager. - # In principle, the internal cache of quantities should be cleared not - # each time the state of the robot and/or its derivative changes. This - # is hard to do because there is no way to detect this specifically at - # the time being. However, `_controller_handle` is never called twice - # in the exact same state by the engine, so resetting quantities at the - # beginning of the method should cover most cases. Yet, quantities - # cannot be used reliably in the definition of profile forces because - # they are always updated before the controller gets called, no matter - # if either one or the other is time-continuous. Hacking the internal - # dynamics to clear quantities does not address this issue either. - self.quantities.clear() - # Refresh the observation self._observer_handle(t, q, v, sensor_measurements) diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py index e62faa75c..c1cc53caf 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py @@ -19,9 +19,11 @@ Callable, cast) import numpy as np + import gymnasium as gym from gymnasium.core import RenderFrame from gymnasium.envs.registration import EnvSpec +from jiminy_py.dynamics import Trajectory from .interfaces import (DT_EPS, ObsT, @@ -31,6 +33,7 @@ InfoType, EngineObsType, InterfaceJiminyEnv) +from .compositions import AbstractReward from .blocks import BaseControllerBlock, BaseObserverBlock from ..utils import DataNested, is_breakpoint, zeros, build_copyto, copy @@ -80,6 +83,7 @@ def __init__(self, self.robot_state = env.robot_state self.sensor_measurements = env.sensor_measurements self.is_simulation_running = env.is_simulation_running + self.num_steps = env.num_steps # Backup the parent environment self.env = env @@ -253,10 +257,7 @@ def step(self, # type: ignore[override] # user keeps doing more steps nonetheless. reward = float(reward) if not math.isnan(reward): - try: - reward += self.compute_reward(terminated, truncated, info) - except NotImplementedError: - pass + reward += self.compute_reward(terminated, info) return obs, reward, terminated, truncated, info @@ -300,6 +301,119 @@ def close(self) -> None: self.env.close() +class ComposedJiminyEnv( + BasePipelineWrapper[ObsT, ActT, ObsT, ActT], + Generic[ObsT, ActT]): + """Extend an environment, eventually already wrapped, by plugging ad-hoc + reward components and termination conditions, including their accompanying + trajectory database if any. + + This wrappers leaves unchanged the observation and action spaces of the + environment. This can be done by adding observation and/or control blocks + through `ObservedJiminyEnv` and `ControlledJiminyEnv` wrappers. + + .. note:: + This wrapper derives from `BasePipelineWrapper`, and such as, it is + considered as internal unlike `gym.Wrapper`. This means that it will be + taken into account when calling `evaluate` or `play_interactive` on the + wrapped environment. + """ + def __init__(self, + env: InterfaceJiminyEnv[ObsT, ActT], + *, + reward: Optional[AbstractReward] = None, + trajectories: Optional[Dict[str, Trajectory]] = None) -> None: + """ + :param env: Environment to extend, eventually already wrapped. + :param reward: Reward object deriving from `AbstractReward`. It will be + evaluated at each step of the environment and summed up + with one returned by the wrapped environment. This + reward must be already instantiated and associated with + the provided environment. `None` for not considering any + reward. + Optional: `None` by default. + :param trajectories: Set of named trajectories as a dictionary whose + (key, value) pairs are respectively the name of + each trajectory and the trajectory itself. `None` + for not considering any trajectory. + Optional: `None` by default. + """ + # Make sure that the unwrapped environment matches the reward one + assert reward is None or env.unwrapped is reward.env.unwrapped + + # Backup user argument(s) + self.reward = reward + + # Initialize base class + super().__init__(env) + + # Add reference trajectories to all managed quantities if requested + if trajectories is not None: + for name, trajectory in trajectories.items(): + self.env.quantities.add_trajectory(name, trajectory) + + # Bind observation and action of the base environment + assert self.observation_space.contains(self.env.observation) + assert self.action_space.contains(self.env.action) + self.observation = self.env.observation + self.action = self.env.action + + def _initialize_action_space(self) -> None: + """Configure the action space. + + It simply copy the action space of the wrapped environment. + """ + self.action_space = self.env.action_space + + def _initialize_observation_space(self) -> None: + """Configure the observation space. + + It simply copy the observation space of the wrapped environment. + """ + self.observation_space = self.env.observation_space + + def _setup(self) -> None: + """Configure the wrapper. + + In addition to calling the base implementation, it sets the observe + and control update period. + """ + # Call base implementation + super()._setup() + + # Copy observe and control update periods from wrapped environment + self.observe_dt = self.env.observe_dt + self.control_dt = self.env.control_dt + + def refresh_observation(self, measurement: EngineObsType) -> None: + """Compute high-level features based on the current wrapped + environment's observation. + + It simply forwards the observation computed by the wrapped environment + without any processing. + + :param measurement: Low-level measure from the environment to process + to get higher-level observation. + """ + self.env.refresh_observation(measurement) + + def compute_command(self, action: ActT, command: np.ndarray) -> None: + """Compute the motors efforts to apply on the robot. + + It simply forwards the command computed by the wrapped environment + without any processing. + + :param action: High-level target to achieve by means of the command. + :param command: Lower-level command to updated in-place. + """ + self.env.compute_command(action, command) + + def compute_reward(self, terminated: bool, info: InfoType) -> float: + if self.reward is None: + return 0.0 + return self.reward(terminated, info) + + class ObservedJiminyEnv( BasePipelineWrapper[NestedObsT, ActT, BaseObsT, ActT], Generic[NestedObsT, ActT, BaseObsT]): @@ -366,8 +480,8 @@ def __init__(self, # Make sure that the environment is either some `ObservedJiminyEnv` or # `ControlledJiminyEnv` block, or the base environment directly. - if isinstance(env, BasePipelineWrapper) and not isinstance( - env, (ObservedJiminyEnv, ControlledJiminyEnv)): + if isinstance(env, BasePipelineWrapper) and not isinstance(env, ( + ObservedJiminyEnv, ControlledJiminyEnv, ComposedJiminyEnv)): raise TypeError( "Observers can only be added on top of another observer, " "controller, or a base environment itself.") @@ -588,8 +702,8 @@ def __init__(self, # Make sure that the environment is either some `ObservedJiminyEnv` or # `ControlledJiminyEnv` block, or the base environment directly. - if isinstance(env, BasePipelineWrapper) and not isinstance( - env, (ObservedJiminyEnv, ControlledJiminyEnv)): + if isinstance(env, BasePipelineWrapper) and not isinstance(env, ( + ObservedJiminyEnv, ControlledJiminyEnv, ComposedJiminyEnv)): raise TypeError( "Controllers can only be added on top of another observer, " "controller, or a base environment itself.") @@ -736,11 +850,8 @@ def compute_command(self, action: ActT, command: np.ndarray) -> None: # the right period. self.env.compute_command(self.env.action, command) - def compute_reward(self, - terminated: bool, - truncated: bool, - info: InfoType) -> float: - return self.controller.compute_reward(terminated, truncated, info) + def compute_reward(self, terminated: bool, info: InfoType) -> float: + return self.controller.compute_reward(terminated, info) class BaseTransformObservation( @@ -762,8 +873,8 @@ class BaseTransformObservation( .. note:: This wrapper derives from `BasePipelineWrapper`, and such as, it is - considered as internal unlike `gym.Wrapper`. This means that it will be - taken into account calling `evaluate` or `play_interactive` on the + considered internal unlike `gym.Wrapper`. This means that it will be + taken into account when calling `evaluate` or `play_interactive` on the wrapped environment. """ def __init__(self, env: InterfaceJiminyEnv[ObsT, ActT]) -> None: @@ -869,8 +980,8 @@ class BaseTransformAction( .. note:: This wrapper derives from `BasePipelineWrapper`, and such as, it is - considered as internal unlike `gym.Wrapper`. This means that it will be - taken into account calling `evaluate` or `play_interactive` on the + considered internal unlike `gym.Wrapper`. This means that it will be + taken into account when calling `evaluate` or `play_interactive` on the wrapped environment. """ def __init__(self, env: InterfaceJiminyEnv[ObsT, ActT]) -> None: diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py b/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py new file mode 100644 index 000000000..065562423 --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py @@ -0,0 +1,968 @@ +"""This module promotes quantities as first-class objects. + +Defining quantities this way allows for standardization of common intermediary +metrics for computation of reward components and and termination conditions, eg +the center of pressure or the average spatial velocity of a frame. Overall, it +greatly reduces code duplication and bugs. + +Apart from that, it offers a dedicated quantity manager that is responsible for +optimizing the computation path, which is expected to significantly increase +the step collection throughput. This speedup is achieved by caching already +computed that did not changed since then, computing redundant intermediary +quantities only once per step, and gathering similar quantities in a large +batch to leverage vectorization of math instructions. +""" +import re +import weakref +from enum import Enum +from weakref import ReferenceType +from abc import ABC, abstractmethod +from collections import OrderedDict +from collections.abc import MutableSet +from dataclasses import dataclass, replace +from functools import partial, wraps +from typing import ( + Any, Dict, List, Optional, Tuple, Generic, TypeVar, Type, Iterator, + Callable, Literal, cast) + +import numpy as np + +import jiminy_py.core as jiminy +from jiminy_py.core import ( # pylint: disable=no-name-in-module + multi_array_copyto) +from jiminy_py.dynamics import State, Trajectory, update_quantities +import pinocchio as pin + +from .interfaces import InterfaceJiminyEnv + + +ValueT = TypeVar('ValueT') + + +class WeakMutableCollection(MutableSet, Generic[ValueT]): + """Mutable unordered list container storing weak reference to objects. + Elements will be discarded when no strong reference to the value exists + anymore, and a user-specified callback will be triggered if any. + + Internally, it is implemented as a set for which uniqueness is + characterized by identity instead of equality operator. + """ + def __init__(self, callback: Optional[Callable[[ + "WeakMutableCollection[ValueT]", ReferenceType + ], None]] = None) -> None: + """ + :param callback: Callback that will be triggered every time an element + is discarded from the container. + Optional: None by default. + """ + self._callback = callback + self._ref_list: List[ReferenceType] = [] + + def __callback__(self, ref: ReferenceType) -> None: + """Internal method that will be called every time an element must be + discarded from the containers, either because it was requested by the + user or because no strong reference to the value exists anymore. + + If a callback has been specified by the user, it will be triggered + after removing the weak reference from the container. + """ + self._ref_list.remove(ref) + if self._callback is not None: + self._callback(self, ref) + + def __contains__(self, obj: Any) -> bool: + """Dunder method to check if a weak reference to a given object is + already stored in the container, which is characterized by identity + instead of equality operator. + + :param obj: Object to look for in the container. + """ + return any(ref() is obj for ref in self._ref_list) + + def __iter__(self) -> Iterator[ValueT]: + """Dunder method that returns an iterator over the objects of the + container for which a string reference still exist. + """ + for ref in self._ref_list: + obj = ref() + if obj is not None: + yield obj + + def __len__(self) -> int: + """Dunder method that returns the length of the container. + """ + return len(self._ref_list) + + def add(self, value: ValueT) -> None: + """Add a new element to the container if not already contained. + + This has no effect if the element is already present. + + :param obj: Object to add to the container. + """ + if value not in self: + self._ref_list.append(weakref.ref(value, self.__callback__)) + + def discard(self, value: ValueT) -> None: + """Remove an element from the container if stored in it. + + This method does not raise an exception when the element is missing. + + :param obj: Object to remove from the container. + """ + if value not in self: + self.__callback__(weakref.ref(value)) + + +class SharedCache(Generic[ValueT]): + """Basic thread local shared cache. + + Its API mimics `std::optional` from the Standard C++ library. All it does + is encapsulating any Python object as a mutable variable, plus exposing a + simple mechanism for keeping track of all "owners" of the cache. + + .. warning:: + This implementation is not thread safe. + """ + + owners: WeakMutableCollection["InterfaceQuantity[ValueT]"] + """Owners of the shared buffer, ie quantities relying on it to store the + result of their evaluation. This information may be useful for determining + the most efficient computation path overall. + + .. note:: + Quantities must add themselves to this list when passing them a shared + cache instance. + + .. note:: + Internally, it stores weak references to avoid holding alive quantities + that could be garbage collected otherwise. `WeakSet` cannot be used + because owners are objects all having the same hash, eg "identical" + quantities. + """ + + def __init__(self) -> None: + """ + """ + # Cached value if any + self._value: Optional[ValueT] = None + + # Whether a value is stored in cached + self._has_value: bool = False + + # Initialize "owners" of the shared buffer. + # Define callback to reset part of the computation graph whenever a + # quantity owning the cache gets garbage collected, namely all + # quantities that may assume at some point the existence of this + # deleted owner to find the adjust their computation path. + def _callback(self: WeakMutableCollection["InterfaceQuantity"], + ref: ReferenceType # pylint: disable=unused-argument + ) -> None: + for owner in self: + while owner.parent is not None: + owner = owner.parent + owner.reset(reset_tracking=True) + + self.owners = WeakMutableCollection(_callback) + + @property + def has_value(self) -> bool: + """Whether a value is stored in cache. + """ + return self._has_value + + def reset(self, *, ignore_auto_refresh: bool = False) -> None: + """Clear value stored in cache if any. + """ + # Clear cache + self._value = None + self._has_value = False + + # Refresh automatically if any cache owner requested it and not ignored + if not ignore_auto_refresh: + for owner in self.owners: + if owner.auto_refresh: + owner.get() + break + + def set(self, value: ValueT) -> None: + """Set value in cache, silently overriding the existing value if any. + + .. warning: + Beware the value is stored by reference for efficiency. It is up to + the user to copy it if necessary. + + :param value: Value to store in cache. + """ + self._value = value + self._has_value = True + + def get(self) -> ValueT: + """Return cached value if any, otherwise raises an exception. + """ + if self._has_value: + return cast(ValueT, self._value) + raise ValueError( + "No value has been stored. Please call 'set' before 'get'.") + + +class InterfaceQuantity(ABC, Generic[ValueT]): + """Interface for generic quantities involved observer-controller blocks, + reward components or termination conditions. + + .. note:: + Quantities are meant to be managed automatically via `QuantityManager`. + Dealing with quantities manually is error-prone, and as such, is + strongly discourage. Nonetheless, power-user that understand the risks + are allowed to do it. + + .. warning:: + Mutual dependency between quantities is disallowed. + + .. warning:: + The user is responsible for implementing the dunder methods `__eq__` + and `__hash__` that characterize identical quantities. This property is + used internally by `QuantityManager` to synchronize cache between + them. It is advised to use decorator `@dataclass(unsafe_hash=True)` for + convenience, but it can also be done manually. + """ + + requirements: Dict[str, "InterfaceQuantity"] + """Intermediary quantities on which this quantity may rely on for its + evaluation at some point, depending on the optimal computation path at + runtime. There values will be exposed to the user as usual properties. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional["InterfaceQuantity"], + requirements: Dict[str, "QuantityCreator"], + *, + auto_refresh: bool = False) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param requirements: Intermediary quantities on which this quantity + depends for its evaluation, as a dictionary + whose keys are tuple gathering their respective + class and all their constructor keyword-arguments + except environment 'env' and parent 'parent. + :param auto_refresh: Whether this quantity must be refreshed + automatically as soon as its shared cache has been + cleared if specified, otherwise this does nothing. + """ + # Backup some of user argument(s) + self.env = env + self.parent = parent + self.auto_refresh = auto_refresh + + # Make sure that all requirement names would be valid as property + requirement_names = requirements.keys() + if any(re.match('[^A-Za-z0-9_]', name) for name in requirement_names): + raise ValueError("The name of all quantity requirements should be " + "ASCII alphanumeric characters plus underscore.") + + # Instantiate intermediary quantities if any + self.requirements: Dict[str, InterfaceQuantity] = { + name: cls(env, self, **kwargs) + for name, (cls, kwargs) in requirements.items()} + + # Shared cache handling + self._cache: Optional[SharedCache[ValueT]] = None + self.has_cache = False + + # Track whether the quantity has been called since previous reset + self._is_active = False + + # Whether the quantity must be re-initialized + self._is_initialized: bool = False + + # Add getter dynamically for user-specified intermediary quantities. + # This approach is hacky but much faster than any of other official + # approach, ie implementing custom a `__getattribute__` method or even + # worst a custom `__getattr__` method. + def get_value(name: str, quantity: InterfaceQuantity) -> Any: + return quantity.requirements[name].get() + + for name in requirement_names: + setattr(type(self), name, property(partial(get_value, name))) + + def __getattr__(self, name: str) -> Any: + """Get access to intermediary quantities as first-class properties, + without having to do it through `requirements`. + + .. warning:: + Accessing quantities this way is convenient, but unfortunately + much slower than do it through `requirements` manually. As a + result, this approach is mainly intended for ease of use while + prototyping. + + :param name: Name of the requested quantity. + """ + try: + return self.__getattribute__('requirements')[name].get() + except KeyError as e: + raise AttributeError( + f"'{type(self)}' object has no attribute '{name}'") from e + + def __dir__(self) -> List[str]: + """Attribute lookup. + + It is mainly used by autocomplete feature of Ipython. It is overloaded + to get consistent autocompletion wrt `getattr`. + """ + return [*super().__dir__(), *self.requirements.keys()] + + @property + def cache(self) -> SharedCache[ValueT]: + """Get shared cache if available, otherwise raises an exception. + + .. warning:: + This method is not meant to be overloaded. + """ + if not self.has_cache: + raise RuntimeError( + "No shared cache has been set for this quantity. Make sure it " + "is managed by some `QuantityManager` instance.") + return cast(SharedCache[ValueT], self._cache) + + @cache.setter + def cache(self, cache: Optional[SharedCache[ValueT]]) -> None: + """Set optional cache variable. When specified, it is used to store + evaluated quantity and retrieve its value later one. + + .. warning:: + Value is stored by reference for efficiency. It is up to the user + to the copy to retain its current value for later use if necessary. + + .. note:: + One may overload this method to encapsulate the cache in a custom + wrapper with specialized 'get' and 'set' methods before passing it + to the base implementation. For instance, to enforce copy of the + cached value before returning it. + """ + # Withdraw this quantity from the owners of its current cache if any + if self._cache is not None: + self._cache.owners.discard(self) + + # Declare this quantity as owner of the cache if specified + if cache is not None: + cache.owners.add(self) + + # Update internal cache attribute + self._cache = cache + self.has_cache = cache is not None + + def is_active(self, any_cache_owner: bool = False) -> bool: + """Whether this quantity is considered active, namely `initialize` has + been called at least once since previous tracking reset. + + :param any_owner: False to check only if this exact instance is active, + True if any of the identical quantities (sharing the + same cache) is considered sufficient. + Optional: False by default. + """ + if not any_cache_owner or self._cache is None: + return self._is_active + return any(owner._is_active for owner in self._cache.owners) + + def get(self) -> ValueT: + """Get cached value of requested quantity if available, otherwise + evaluate it and store it in cache. + + This quantity is considered active as soon as this method has been + called at least once since previous tracking reset. The method + `is_active` will be return true even before calling `initialize`. + + .. warning:: + This method is not meant to be overloaded. + """ + # Get value in cache if available. + # Note that direct access to internal `_value` attribute is preferred + # over the public API `get` for speedup. The same cannot be done for + # `has_value` as it would prevent mocking it during running unit tests + # or benchmarks. + if (self.has_cache and + self._cache.has_value): # type: ignore[union-attr] + self._is_active = True + return self._cache._value # type: ignore[union-attr,return-value] + + # Evaluate quantity + try: + if not self._is_initialized: + self.initialize() + assert (self._is_initialized and + self._is_active) # type: ignore[unreachable] + value = self.refresh() + except RecursionError as e: + raise LookupError( + "Mutual dependency between quantities is disallowed.") from e + + # Return value after storing it in shared cache if available + if self.has_cache: + self._cache.set(value) # type: ignore[union-attr] + return value + + def reset(self, reset_tracking: bool = False) -> None: + """Consider that the quantity must be re-initialized before being + evaluated once again. + + If shared cache is available, then it will be cleared and all identity + quantities will jointly be reset. + + .. note:: + This method must be called right before performing agent steps, + otherwise this quantity will not be refreshed if it was evaluated + previously. + + .. warning:: + This method is not meant to be overloaded. + + :param reset_tracking: Do not consider this quantity as active anymore + until the `get` method gets called once again. + Optional: False by default. + """ + # Make sure that auto-refresh can be honored + if self.auto_refresh and not self.has_cache: + raise RuntimeError( + "Automatic refresh enabled but no shared cache available. " + "Please add one before calling this method.") + + # No longer consider this exact instance as initialized + self._is_initialized = False + + # No longer consider this exact instance as active if requested + if reset_tracking: + self._is_active = False + + # Reset all requirements first + for quantity in self.requirements.values(): + quantity.reset(reset_tracking) + + # More work must to be done if shared cache is available and has value + if self.has_cache: + # Early return if shared cache has no value + if not self.cache.has_value: + return + + # Invalidate cache before looping over all identical properties. + # Note that auto-refresh must be ignored to avoid infinite loop. + self.cache.reset(ignore_auto_refresh=True) + + # Reset all identical quantities + for owner in self.cache.owners: + owner.reset() + + # Reset shared cache one last time but without ignore auto refresh + self.cache.reset() + + def initialize(self) -> None: + """Initialize internal buffers. + + This is typically useful to refresh shared memory proxies or to + re-initialize pre-allocated buffers. + + .. warning:: + Intermediary quantities 'requirements' are NOT initialized + automatically because they can be initialized lazily in most cases, + or are optional depending on the most efficient computation path at + run-time. It is up to the developer implementing quantities to take + care of it. + + .. note:: + This method must be called before starting a new episode. + + .. note:: + Lazy-initialization is used for efficiency, ie `initialize` will be + called before the first time `refresh` has to be called, which may + never be the case if cache is shared between multiple identical + instances of the same quantity. + """ + # The quantity is now considered initialized and active unconditionally + self._is_initialized = True + self._is_active = True + + @abstractmethod + def refresh(self) -> ValueT: + """Evaluate this quantity based on the agent state at the end of the + current agent step. + """ + + +QuantityCreator = Tuple[Type[InterfaceQuantity[ValueT]], Dict[str, Any]] + + +class QuantityEvalMode(Enum): + """Specify on which state to evaluate a given quantity. + """ + + TRUE = 0 + """Current state of the environment. + """ + + REFERENCE = 1 + """State of the reference trajectory at the current simulation time. + """ + + +@dataclass(unsafe_hash=True) +class AbstractQuantity(InterfaceQuantity, Generic[ValueT]): + """Base class for generic quantities involved observer-controller blocks, + reward components or termination conditions. + + .. note:: + A dataset of trajectories made available through `self.trajectories`. + The latter is synchronized because all quantities as long as shared + cached is available. Since the dataset is initially empty by default, + using `QuantityEvalMode.REFERENCE` evaluation mode requires manually + adding at least one trajectory to the dataset and selecting it. + + .. seealso:: + See `InterfaceQuantity` documentation for details. + """ + + mode: QuantityEvalMode + """Specify on which state to evaluate this quantity. See `Mode` + documentation for details about each mode. + + .. warning:: + Mode `REFERENCE` requires a reference trajectory to be selected + manually prior to evaluating this quantity for the first time. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + requirements: Dict[str, "QuantityCreator"], + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE, + auto_refresh: bool = False) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param requirements: Intermediary quantities on which this quantity + depends for its evaluation, as a dictionary + whose keys are tuple gathering their respective + class and all their constructor keyword-arguments + except environment 'env' and parent 'parent. + :param mode: Desired mode of evaluation for this quantity. If mode is + set to `QuantityEvalMode.TRUE`, then current simulation + state will be used in dynamics computations. If mode is + set to `QuantityEvalMode.REFERENCE`, then at the state of + some reference trajectory at the current simulation time + will be used instead. + :param auto_refresh: Whether this quantity must be refreshed + automatically as soon as its shared cache has been + cleared if specified, otherwise this does nothing. + """ + # Backup user argument(s) + self.mode = mode + + # Make sure that no user-specified requirement is named 'trajectory' + requirement_names = requirements.keys() + if any(name in requirement_names for name in ("state", "trajectory")): + raise ValueError( + "No user-specified requirement can be named 'state' nor " + "'trajectory' as these keys are reserved.") + + # Add state quantity as requirement + requirements["state"] = (StateQuantity, dict(mode=mode)) + + # Call base implementation + super().__init__(env, parent, requirements, auto_refresh=auto_refresh) + + # Add trajectory quantity proxy + trajectory = self.requirements["state"].requirements["trajectory"] + assert isinstance(trajectory, DatasetTrajectoryQuantity) + self.trajectory = trajectory + + # Robot for which the quantity must be evaluated + self.robot = jiminy.Robot() + self.pinocchio_model = pin.Model() + self.pinocchio_data = self.pinocchio_model.createData() + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Refresh robot proxy + state = self.requirements["state"] + assert isinstance(state, StateQuantity) + self.robot = state.robot + self.pinocchio_model = state.pinocchio_model + self.pinocchio_data = state.pinocchio_data + + +def sync(fun: Callable[..., None]) -> Callable[..., None]: + """Wrap any `InterfaceQuantity` instance method to forward call to all + co-owners of the same shared cache. + + This wrapper is useful to keep all identical instances of the same quantity + in sync. + """ + @wraps(fun) + def fun_safe(self: InterfaceQuantity, *args: Any, **kwargs: Any) -> None: + # Hijack instance for adding private an attribute tracking whether its + # internal state went out-of-sync between identical instances. + # Note that a local variable cannot be used because all synched methods + # must shared the same tracking state variable. Otherwise, one method + # may be flagged out-of-sync but not the others. + if not hasattr(self, "__is_synched__"): + self.__is_synched__ = self.has_cache # type: ignore[attr-defined] + + # Check if quantity has cache but is already out-of-sync. + # Raise exception if it now has cache while it was not the case before. + must_sync = self.has_cache and len(self.cache.owners) > 1 + if not self.__is_synched__ and must_sync: + raise RuntimeError( + "This quantity went out-of-sync. Make sure that no synched " + "method is called priori to setting shared cache.") + self.__is_synched__ = self.has_cache # type: ignore[attr-defined] + + # Call instance method on the original caller first + fun(self, *args, **kwargs) + + # Call instance method on all other co-owners of shared cache if any + if self.has_cache: + cls = type(self) + for owner in self.cache.owners: + if owner is not self: + assert isinstance(owner, cls) + value = fun(owner, *args, **kwargs) + if value is not None: + raise NotImplementedError( + "Instance methods that does not return `None` are " + "not supported.") + + return fun_safe + + +@dataclass(unsafe_hash=True) +class DatasetTrajectoryQuantity(InterfaceQuantity[State]): + """This class manages a database of trajectories. + + The database is empty by default. Trajectories must be added or discarded + manually. Only one trajectory can be selected at once. Once a trajectory + has been selecting, its state at the current simulation can be easily + retrieved. + + This class does not require to only adding trajectories for which all + attributes of the underlying state sequence have been specified. Missing + attributes of a trajectory will also be missing from the retrieved state. + It is the responsible of the user to make sure all cases are properly + handled if needed. + + All instances of this quantity sharing the same cache are synchronized, + which means that adding, discarding, or selecting a trajectory on any of + them would propagate on all the others. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + ) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + """ + # Call base implementation + super().__init__(env, parent, requirements={}, auto_refresh=False) + + # Ordered set of named reference trajectories as a dictionary + self.registry: OrderedDict[str, Trajectory] = OrderedDict() + + # Name of the trajectory that is currently selected + self._name = "" + + # Selected trajectory if any + self._trajectory: Optional[Trajectory] = None + + # Specifies how to deal with query time that are out-of-bounds + self._mode: Literal['raise', 'wrap', 'clip'] = 'raise' + + @property + def trajectory(self) -> Trajectory: + """Trajectory that is currently selected if any, raises an exception + otherwise. + """ + # Make sure that a trajectory has been selected + if self._trajectory is None: + raise RuntimeError("No trajectory has been selected.") + + # Return selected trajectory + return self._trajectory + + @property + def robot(self) -> jiminy.Robot: + """Robot associated with the selected trajectory. + """ + return self.trajectory.robot + + @property + def use_theoretical_model(self) -> bool: + """Whether the selected trajectory is associated with the theoretical + dynamical model or extended simulation model of the robot. + """ + return self.trajectory.use_theoretical_model + + @sync + def _add(self, name: str, trajectory: Trajectory) -> None: + """Add a trajectory to local internal registry only without performing + any validity check. + + .. warning:: + This method is used internally by `add` method. It is not meant to + be called manually. + + :param name: Desired name of the trajectory. + :param trajectory: Trajectory instance to register. + """ + self.registry[name] = trajectory + + def add(self, name: str, trajectory: Trajectory) -> None: + """Jointly add a trajectory to the local internal registry of all + instances sharing the same cache as this quantity. + + :param name: Desired name of the trajectory. It must be unique. If a + trajectory with the exact same name already exists, then + it must be discarded first, so as to prevent silently + overwriting it by mistake. + :param trajectory: Trajectory instance to register. + """ + # Make sure that no trajectory with the exact same name already exists + if name in self.registry: + raise KeyError( + "A trajectory with the exact same name already exists. Please " + "delete it first before adding a new one.") + + # Allocate new dummy robot to avoid altering the simulation one + if trajectory.robot is self.env.robot: + trajectory = replace(trajectory, robot=trajectory.robot.copy()) + + # Add the same post-processed trajectory to all identical instances. + # Note that `add` must be splitted in two methods. A first part that + # applies some trajectory post-processing only once, and a second part + # that adds the post-processed trajectory to all identical quantities + # at once. It is absolutely essential to proceed this way, because it + # guarantees that the underlying trajectories are all references to the + # same memory, including `pinocchio_data`. This means that calling + # `update_quantities` will perform the update for all of them at once. + # Consequently, kinematics and dynamics quantities of all `State` + # instances will be up-to-date as long as `refresh` is called once for + # a given evaluation mode. + self._add(name, trajectory) + + @sync + def discard(self, name: str) -> None: + """Jointly remove a trajectory from the local internal registry of all + instances sharing the same cache as this quantity. + + :param name: Name of the trajectory to discard. + """ + # Un-select trajectory if it corresponds to the discarded one + if self._name == name: + self._trajectory = None + self._name = "" + + # Delete trajectory for global registry + del self.registry[name] + + @sync + def select(self, + name: str, + mode: Literal['raise', 'wrap', 'clip'] = 'raise') -> None: + """Jointly select a trajectory in the internal registry of all + instances sharing the same cache as this quantity. + + :param name: Name of the trajectory to discard. + :param mode: Specifies how to deal with query time of are out of the + time interval of the trajectory. See `Trajectory.get` + documentation for details. + """ + # Make sure that at least one trajectory has been specified + if not self.registry: + raise ValueError("Cannot select trajectory on a empty dataset.") + + # Select the desired trajectory for all identical instances + self._trajectory = self.registry[name] + self._name = name + + # Backup user-specified mode + self._mode = mode + + # Un-initialize quantity when the selected trajectory changes + self.reset(reset_tracking=False) + + @property + def name(self) -> str: + """Name of the trajectory that is currently selected. + """ + return self._name + + @InterfaceQuantity.cache.setter # type: ignore[attr-defined] + def cache(self, cache: Optional[SharedCache[ValueT]]) -> None: + # Get existing registry if any and making sure not already out-of-sync + owner: Optional[InterfaceQuantity] = None + if cache is not None and cache.owners: + owner = next(iter(cache.owners)) + assert isinstance(owner, DatasetTrajectoryQuantity) + if self._trajectory: + raise RuntimeError( + "Trajectory dataset not empty. Impossible to add a shared " + "cache already having owners.") + + # Call base implementation + InterfaceQuantity.cache.fset(self, cache) # type: ignore[attr-defined] + + # Catch-up synchronization + if owner: + self.registry = owner.registry + if owner._trajectory is not None: + self.select(owner._name, owner._mode) + + def refresh(self) -> State: + """Compute state of selected trajectory at current simulation time. + """ + return self.trajectory.get(self.env.stepper_state.t, self._mode) + + +@dataclass(unsafe_hash=True) +class StateQuantity(InterfaceQuantity[State]): + """State to consider when evaluating any quantity deriving from + `AbstractQuantity` using the same evaluation mode as this instance. + + This quantity is refreshed automatically no matter what. This guarantees + that all low-level kinematics and dynamics quantities that can be computed + from the current state are up-to-date. More specifically, every quantities + would be up-to-date if the evaluation mode is `QuantityEvalMode.TRUE`, + while it would depends on the information available on the selected + trajectory if the evaluation mode is `QuantityEvalMode.REFERENCE`. See + `update_quantities` documentation for details. + """ + + mode: QuantityEvalMode + """Specify on which state to evaluate this quantity. See `Mode` + documentation for details about each mode. + + .. warning:: + Mode `REFERENCE` requires a reference trajectory to be selected + manually prior to evaluating this quantity for the first time. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param mode: Desired mode of evaluation for this quantity. If mode is + set to `QuantityEvalMode.TRUE`, then current simulation + state will be used in dynamics computations. If mode is + set to `QuantityEvalMode.REFERENCE`, then at the state of + some reference trajectory at the current simulation time + will be used instead. + """ + # Backup user argument(s) + self.mode = mode + + # Call base implementation + super().__init__(env, parent, requirements={}, auto_refresh=True) + + # Create empty trajectory database, manually added as a requirement. + # Note that it must be done after base initialization, otherwise a + # getter will be added for it as first-class property. + self.trajectory = DatasetTrajectoryQuantity(env, self) + self.requirements["trajectory"] = self.trajectory + + # Robot for which the quantity must be evaluated + self.robot = env.robot + self.pinocchio_model = env.robot.pinocchio_model + self.pinocchio_data = env.robot.pinocchio_data + + # State for which the quantity must be evaluated + self._f_external_slices: Tuple[np.ndarray, ...] = () + self._f_external_list: Tuple[np.ndarray, ...] = () + self.state = State(t=np.nan, q=np.array([])) + + def initialize(self) -> None: + # Refresh robot and pinocchio proxies for co-owners of shared cache. + # Note that automatic refresh is not sufficient to guarantee that + # `initialize` will be called unconditionally, because it will be + # skipped if a value is already stored in cache. As a result, it is + # necessary to synchronize calls to this method between co-owners of + # the shared cache manually, so that it will be called by the first + # instance to found the cache empty. Only the necessary bits are + # synchronized instead of the whole method, to avoid messing up with + # computation graph tracking. + owners = self.cache.owners if self.has_cache else (self,) + for owner in owners: + assert isinstance(owner, StateQuantity) + if owner._is_initialized: + continue + if owner.mode == QuantityEvalMode.TRUE: + owner.robot = owner.env.robot + use_theoretical_model = False + else: + owner.robot = owner.trajectory.robot + use_theoretical_model = owner.trajectory.use_theoretical_model + if use_theoretical_model: + owner.pinocchio_model = owner.robot.pinocchio_model_th + owner.pinocchio_data = owner.robot.pinocchio_data_th + else: + owner.pinocchio_model = owner.robot.pinocchio_model + owner.pinocchio_data = owner.robot.pinocchio_data + + # Call base implementation. + # Thz quantity will be considered initialized and active at this point. + super().initialize() + + # State for which the quantity must be evaluated + if self.mode == QuantityEvalMode.TRUE: + self._f_external_list = tuple( + f_ext.vector for f_ext in self.env.robot_state.f_external) + if self._f_external_list: + f_external_batch = np.stack(self._f_external_list, axis=0) + else: + f_external_batch = np.array([]) + self.state = State( + self.env.stepper_state.t, + self.env.robot_state.q, + self.env.robot_state.v, + self.env.robot_state.a, + self.env.robot_state.u, + self.env.robot_state.command, + f_external_batch) + self._f_external_slices = tuple(f_external_batch) + + def refresh(self) -> State: + """Compute the current state depending on the mode of evaluation, and + make sure that kinematics and dynamics quantities are up-to-date. + """ + # Update state at which the quantity must be evaluated + if self.mode == QuantityEvalMode.TRUE: + multi_array_copyto(self._f_external_slices, self._f_external_list) + else: + self.state = self.trajectory.get() + + # Update all the dynamical quantities that can be given available data + if self.mode == QuantityEvalMode.REFERENCE: + update_quantities( + self.robot, + self.state.q, + self.state.v, + self.state.a, + update_physics=True, + update_centroidal=True, + update_energy=True, + update_jacobian=False, + update_collisions=True, + use_theoretical_model=self.trajectory.use_theoretical_model) + + # Return state + return self.state diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/quantity.py b/python/gym_jiminy/common/gym_jiminy/common/bases/quantity.py deleted file mode 100644 index e2dbefc59..000000000 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/quantity.py +++ /dev/null @@ -1,454 +0,0 @@ -"""This module promotes quantities as first-class objects. - -Defining quantities this way allows for standardization of common intermediary -metrics for computation of reward components and and termination conditions, eg -the center of pressure or the average spatial velocity of a frame. Overall, it -greatly reduces code duplication and bugs. - -Apart from that, it offers a dedicated quantity manager that is responsible for -optimizing the computation path, which is expected to significantly increase -the step collection throughput. This speedup is achieved by caching already -computed that did not changed since then, computing redundant intermediary -quantities only once per step, and gathering similar quantities in a large -batch to leverage vectorization of math instructions. -""" -import weakref -from weakref import ReferenceType -from abc import ABC, abstractmethod -from collections.abc import MutableSet -from functools import partial -from typing import ( - Any, Dict, List, Optional, Tuple, Generic, TypeVar, Type, Iterator, - Callable, cast) - -from .interfaces import InterfaceJiminyEnv - - -ValueT = TypeVar('ValueT') - -QuantityCreator = Tuple[Type["AbstractQuantity"], Dict[str, Any]] - - -class WeakMutableCollection(MutableSet, Generic[ValueT]): - """Mutable unordered list container storing weak reference to objects. - Elements will be discarded when no strong reference to the value exists - anymore, and a user-specified callback will be triggered if any. - - Internally, it is implemented as a set for which uniqueness is - characterized by identity instead of equality operator. - """ - def __init__(self, callback: Optional[Callable[[ - "WeakMutableCollection[ValueT]", ReferenceType - ], None]] = None) -> None: - """ - :param callback: Callback that will be triggered every time an element - is discarded from the container. - Optional: None by default. - """ - self._callback = callback - self._ref_list: List[ReferenceType] = [] - - def __callback__(self, ref: ReferenceType) -> None: - """Internal method that will be called every time an element must be - discarded from the containers, either because it was requested by the - user or because no strong reference to the value exists anymore. - - If a callback has been specified by the user, it will be triggered - after removing the weak reference from the container. - """ - self._ref_list.remove(ref) - if self._callback is not None: - self._callback(self, ref) - - def __contains__(self, obj: Any) -> bool: - """Dunder method to check if a weak reference to a given object is - already stored in the container, which is characterized by identity - instead of equality operator. - - :param obj: Object to look for in the container. - """ - return any(ref() is obj for ref in self._ref_list) - - def __iter__(self) -> Iterator[ValueT]: - """Dunder method that returns an iterator over the objects of the - container for which a string reference still exist. - """ - for ref in self._ref_list: - obj = ref() - if obj is not None: - yield obj - - def __len__(self) -> int: - """Dunder method that returns the length of the container. - """ - return len(self._ref_list) - - def add(self, value: ValueT) -> None: - """Add a new element to the container if not already contained. - - This has no effect if the element is already present. - - :param obj: Object to add to the container. - """ - if value not in self: - self._ref_list.append(weakref.ref(value, self.__callback__)) - - def discard(self, value: ValueT) -> None: - """Remove an element from the container if stored in it. - - This method does not raise an exception when the element is missing. - - :param obj: Object to remove from the container. - """ - if value not in self: - self.__callback__(weakref.ref(value)) - - -class SharedCache(Generic[ValueT]): - """Basic thread local shared cache. - - Its API mimics `std::optional` from the Standard C++ library. All it does - is encapsulating any Python object as a mutable variable, plus exposing a - simple mechanism for keeping track of all "owners" of the cache. - - .. warning:: - This implementation is not thread safe. - """ - - owners: WeakMutableCollection["AbstractQuantity"] - """Owners of the shared buffer, ie quantities relying on it to store the - result of their evaluation. This information may be useful for determining - the most efficient computation path overall. - - .. note:: - Quantities must add themselves to this list when passing them a shared - cache instance. - - .. note:: - Internally, it stores weak references to avoid holding alive quantities - that could be garbage collected otherwise. `WeakSet` cannot be used - because owners are objects all having the same hash, eg "identical" - quantities. - """ - - def __init__(self) -> None: - """ - """ - # Cached value if any - self._value: Optional[ValueT] = None - - # Whether a value is stored in cached - self._has_value: bool = False - - # Initialize "owners" of the shared buffer. - # Define callback to reset part of the computation graph whenever a - # quantity owning the cache gets garbage collected, namely all - # quantities that may assume at some point the existence of this - # deleted owner to find the adjust their computation path. - def _callback(self: WeakMutableCollection["AbstractQuantity"], - ref: ReferenceType # pylint: disable=unused-argument - ) -> None: - for owner in self: - while owner.parent is not None: - owner = owner.parent - owner.reset(reset_tracking=True) - - self.owners = WeakMutableCollection(_callback) - - @property - def has_value(self) -> bool: - """Whether a value is stored in cache. - """ - return self._has_value - - def reset(self) -> None: - """Clear value stored in cache if any. - """ - self._value = None - self._has_value = False - - def set(self, value: ValueT) -> None: - """Set value in cache, silently overriding the existing value if any. - - .. warning: - Beware the value is stored by reference for efficiency. It is up to - the user to copy it if necessary. - - :param value: Value to store in cache. - """ - self._value = value - self._has_value = True - - def get(self) -> ValueT: - """Return cached value if any, otherwise raises an exception. - """ - if self._has_value: - return cast(ValueT, self._value) - raise ValueError( - "No value has been stored. Please call 'set' before 'get'.") - - -class AbstractQuantity(ABC, Generic[ValueT]): - """Interface for quantities that involved in reward or termination - conditions evaluation. - - .. note:: - Quantities are meant to be managed automatically via `QuantityManager`. - Dealing with quantities manually is error-prone, and as such, is - strongly discourage. Nonetheless, power-user that understand the risks - are allowed to do it. - - .. warning:: - Mutual dependency between quantities is disallowed. - - .. warning:: - The user is responsible for implementing the dunder methods `__eq__` - and `__hash__` that characterize identical quantities. This property is - used internally by `QuantityManager` to synchronize cache between - them. It is advised to use decorator `@dataclass(unsafe_hash=True)` for - convenience, but it can also be done manually. - """ - - requirements: Dict[str, "AbstractQuantity"] - """Intermediary quantities on which this quantity may rely on for its - evaluation at some point, depending on the optimal computation path at - runtime. There values will be exposed to the user as usual properties. - """ - - def __init__(self, - env: InterfaceJiminyEnv, - parent: Optional["AbstractQuantity"], - requirements: Dict[str, QuantityCreator]) -> None: - """ - :param env: Base or wrapped jiminy environment. - :param parent: Higher-level quantity from which this quantity is a - requirement if any, `None` otherwise. - :param requirements: Intermediary quantities on which this quantity - depends for its evaluation, as a dictionary - whose keys are tuple gathering their respective - class and all their constructor keyword-arguments - except the environment 'env'. - """ - # Backup some of user argument(s) - self.env = env - self.parent = parent - - # Instantiate intermediary quantities if any - self.requirements: Dict[str, AbstractQuantity] = { - name: cls(env, self, **kwargs) - for name, (cls, kwargs) in requirements.items()} - - # Define some proxies for fast access - self.pinocchio_model = env.robot.pinocchio_model - self.pinocchio_data = env.robot.pinocchio_data - - # Shared cache handling - self._cache: Optional[SharedCache[ValueT]] = None - self._has_cache = False - - # Track whether the quantity has been called since previous reset - self._is_active = False - - # Whether the quantity must be re-initialized - self._is_initialized: bool = False - - # Add getter of all intermediary quantities dynamically. - # This approach is hacky but much faster than any of other official - # approach, ie implementing custom a `__getattribute__` method or even - # worst a custom `__getattr__` method. - def get_value(name: str, quantity: AbstractQuantity) -> Any: - return quantity.requirements[name].get() - - for name in self.requirements.keys(): - setattr(type(self), name, property(partial(get_value, name))) - - def __getattr__(self, name: str) -> Any: - """Get access to intermediary quantities as first-class properties, - without having to do it through `requirements`. - - .. warning:: - Accessing quantities this way is convenient, but unfortunately - much slower than do it through `requirements` manually. As a - result, this approach is mainly intended for ease of use while - prototyping. - - :param name: Name of the requested quantity. - """ - return self.__getattribute__('requirements')[name].get() - - def __dir__(self) -> List[str]: - """Attribute lookup. - - It is mainly used by autocomplete feature of Ipython. It is overloaded - to get consistent autocompletion wrt `getattr`. - """ - return [*super().__dir__(), *self.requirements.keys()] - - @property - def cache(self) -> SharedCache[ValueT]: - """Get shared cache if available, otherwise raises an exception. - - .. warning:: - This method is not meant to be overloaded. - """ - if not self._has_cache: - raise RuntimeError( - "No shared cache has been set for this quantity. Make sure it " - "is managed by some `QuantityManager` instance.") - return cast(SharedCache[ValueT], self._cache) - - @cache.setter - def cache(self, cache: Optional[SharedCache[ValueT]]) -> None: - """Set optional cache variable. When specified, it is used to store - evaluated quantity and retrieve its value later one. - - .. warning:: - Value is stored by reference for efficiency. It is up to the user - to the copy to retain its current value for later use if necessary. - - .. note:: - One may overload this method to encapsulate the cache in a custom - wrapper with specialized 'get' and 'set' methods before passing it - to the base implementation. For instance, to enforce copy of the - cached value before returning it. - """ - # Withdraw this quantity from the owners of its current cache if any - if self._cache is not None: - self._cache.owners.discard(self) - - # Declare this quantity as owner of the cache if specified - if cache is not None: - cache.owners.add(self) - - # Update internal cache attribute - self._cache = cache - self._has_cache = cache is not None - - def is_active(self, any_cache_owner: bool = False) -> bool: - """Whether this quantity is considered active, namely `initialize` has - been called at least once since previous tracking reset. - - :param any_owner: False to check only if this exact instance is active, - True if any of the identical quantities (sharing the - same cache) is considered sufficient. - Optional: False by default. - """ - if not any_cache_owner or self._cache is None: - return self._is_active - return any(owner._is_active for owner in self._cache.owners) - - def get(self) -> ValueT: - """Get cached value of requested quantity if available, otherwise - evaluate it and store it in cache. - - This quantity is considered active as soon as this method has been - called at least once since previous tracking reset. The method - `is_active` will be return true even before calling `initialize`. - - .. warning:: - This method is not meant to be overloaded. - """ - # Get value in cache if available. - # Note that direct access to internal `_value` attribute is preferred - # over the public API `get` for speedup. The same cannot be done for - # `has_value` as it would prevent mocking it during running unit tests - # or benchmarks. - if (self._has_cache and - self._cache.has_value): # type: ignore[union-attr] - self._is_active = True - return self._cache._value # type: ignore[union-attr,return-value] - - # Evaluate quantity - try: - if not self._is_initialized: - self.initialize() - assert (self._is_initialized and - self._is_active) # type: ignore[unreachable] - value = self.refresh() - except RecursionError as e: - raise LookupError( - "Mutual dependency between quantities is disallowed.") from e - - # Return value after storing it in shared cache if available - if self._has_cache: - self._cache.set(value) # type: ignore[union-attr] - return value - - def reset(self, reset_tracking: bool = False) -> None: - """Consider that the quantity must be re-initialized before being - evaluated once again. - - If shared cache is available, then it will be cleared and all identity - quantities will jointly be reset. - - .. note:: - This method must be called right before performing agent steps, - otherwise this quantity will not be refreshed if it was evaluated - previously. - - .. warning:: - This method is not meant to be overloaded. - - :param reset_tracking: Do not consider this quantity as active anymore - until the `get` method gets called once again. - Optional: False by default. - """ - # No longer consider this exact instance as initialized - self._is_initialized = False - - # No longer consider this exact instance as active if requested - if reset_tracking: - self._is_active = False - - # Reset all requirements first - for quantity in self.requirements.values(): - quantity.reset(reset_tracking) - - # More work has to be done if shared cache is available and has value - if self._has_cache: - # Early return if shared cache has no value - if not self.cache.has_value: - return - - # Invalidate cache before looping over all identical properties - self.cache.reset() - - # Reset all identical quantities - for owner in self.cache.owners: - owner.reset() - - def initialize(self) -> None: - """Initialize internal buffers. - - This is typically useful to refresh shared memory proxies or to - re-initialize pre-allocated buffers. - - .. warning:: - Intermediary quantities 'requirements' are NOT initialized - automatically because they can be initialized lazily in most cases, - or are optional depending on the most efficient computation path at - run-time. It is up to the developer implementing quantities to take - care of it. - - .. note:: - This method must be called before starting a new episode. - - .. note:: - Lazy-initialization is used for efficiency, ie `initialize` will be - called before the first time `refresh` has to be called, which may - never be the case if cache is shared between multiple identical - instances of the same quantity. - """ - # Refresh some proxies - self.pinocchio_model = self.env.robot.pinocchio_model - self.pinocchio_data = self.env.robot.pinocchio_data - - # The quantity is now considered initialized and active unconditionally - self._is_initialized = True - self._is_active = True - - @abstractmethod - def refresh(self) -> ValueT: - """Evaluate this quantity based on the agent state at the end of the - current agent step. - """ diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py index 902b7bfd1..cdcac2029 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py @@ -10,7 +10,7 @@ import jiminy_py.core as jiminy from jiminy_py.core import ( # pylint: disable=no-name-in-module - EncoderSensor as encoder, ImuSensor as imu, get_frame_indices) + EncoderSensor, ImuSensor, get_frame_indices) import pinocchio as pin @@ -260,8 +260,10 @@ def get_flexibility_imu_frame_chains( supports = [] for joint_index in leaf_joint_indices: support = [] - while joint_index > 0: + while True: support.append(joint_index) + if joint_index == 0: + break joint_index = parents[joint_index] supports.append(support) @@ -544,16 +546,17 @@ def __init__(self, # Backup some of the user-argument(s) self.ignore_twist = ignore_twist - # Define proxies for fast access - self.pinocchio_model_th = env.robot.pinocchio_model_th.copy() - self.pinocchio_data_th = env.robot.pinocchio_data_th.copy() - # Create flexible dynamic model. # Dummy physical parameters are specified as they have no effect on # kinematic computations. model = jiminy.Model() - model.initialize(env.robot.pinocchio_model_th) + pinocchio_model_th = env.robot.pinocchio_model_th + if env.robot.has_freeflyer: + pinocchio_model_th = pin.buildReducedModel( + pinocchio_model_th, [1], pin.neutral(pinocchio_model_th)) + model.initialize(pinocchio_model_th) options = model.get_options() + options["dynamics"]["enableFlexibility"] = True for frame_name in flex_frame_names: options["dynamics"]["flexibilityConfig"].append( { @@ -565,6 +568,10 @@ def __init__(self, ) model.set_options(options) + # Backup theoretical pinocchio model without floating base + self.pinocchio_model_th = model.pinocchio_model_th.copy() + self.pinocchio_data_th = model.pinocchio_data_th.copy() + # Extract contiguous chains of flexibility and IMU frames for which # computations can be vectorized. It also stores the information of # whether or not the sign of the deformation must be reversed to be @@ -636,34 +643,37 @@ def __init__(self, # Get mapping from IMU frame to index imu_frame_map: Dict[str, int] = {} - for sensor_name in env.robot.sensor_names[imu.type]: - sensor = env.robot.get_sensor(imu.type, sensor_name) - assert isinstance(sensor, imu) + for sensor in env.robot.sensors[ImuSensor.type]: + assert isinstance(sensor, ImuSensor) imu_frame_map[sensor.frame_name] = sensor.index # Make sure that the robot has one encoder per mechanical joint - encoder_sensor_names = env.robot.sensor_names[encoder.type] - if len(encoder_sensor_names) < len(model.mechanical_joint_indices): + encoder_sensors = env.robot.sensors[EncoderSensor.type] + if len(encoder_sensors) < len(model.mechanical_joint_indices): raise ValueError( "The robot must have one encoder per mechanical joints.") # Extract mapping from encoders to theoretical configuration. # Note that revolute unbounded joints are not supported for now. - self.encoder_to_config = [-1 for _ in range(env.robot.nmotors)] - for i, sensor_name in enumerate(encoder_sensor_names): - sensor = env.robot.get_sensor(encoder.type, sensor_name) - assert isinstance(sensor, encoder) - if sensor.joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: + self.encoder_to_position_map = [-1,] * env.robot.nmotors + for sensor in env.robot.sensors[EncoderSensor.type]: + assert isinstance(sensor, EncoderSensor) + joint_index = self.pinocchio_model_th.getJointId(sensor.joint_name) + joint = self.pinocchio_model_th.joints[joint_index] + joint_type = jiminy.get_joint_type(joint) + if joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: raise ValueError( "Revolute unbounded joints are not supported for now.") - encoder_joint = self.pinocchio_model_th.joints[sensor.joint_index] - self.encoder_to_config[i] = encoder_joint.idx_q + self.encoder_to_position_map[sensor.index] = joint.idx_q - # Extract measured joint positions for fast access. + # Extract measured motor / joint positions for fast access. # Note that they will be initialized in `_setup` method. self.encoder_data = np.array([]) - # Buffer storing the theoretical configuration. + # Ratio to translate encoder data to joint side + self.encoder_to_joint_ratio = np.array([]) + + # Buffer storing the theoretical configuration self._q_th = pin.neutral(self.pinocchio_model_th) # Whether the observer has been compiled already @@ -702,8 +712,12 @@ def _setup(self) -> None: # Refresh the theoretical model of the robot. # Even if the robot may change, the theoretical model of the robot is # not supposed to change in a way that would break this observer. - self.pinocchio_model_th = self.env.robot.pinocchio_model_th - self.pinocchio_data_th = self.env.robot.pinocchio_data_th + pinocchio_model_th = self.env.robot.pinocchio_model_th + if self.env.robot.has_freeflyer: + pinocchio_model_th = pin.buildReducedModel( + pinocchio_model_th, [1], pin.neutral(pinocchio_model_th)) + self.pinocchio_model_th = pinocchio_model_th + self.pinocchio_data_th = self.pinocchio_model_th.createData() # Fix initialization of the observation to be valid quaternions self.observation[-1] = 1.0 @@ -722,8 +736,20 @@ def _setup(self) -> None: self._kin_flex_rots.append(kin_flex_rots) self._kin_imu_rots.append(kin_imu_rots) - # Refresh measured motor positions and velocities proxies - self.encoder_data, _ = self.env.sensor_measurements[encoder.type] + # Refresh measured motor position proxy + self.encoder_data, _ = self.env.sensor_measurements[EncoderSensor.type] + + # Refresh mechanical reduction ratio + encoder_to_joint_ratio = [] + for sensor in self.env.robot.sensors[EncoderSensor.type]: + try: + motor = self.env.robot.motors[sensor.motor_index] + motor_options = motor.get_options() + mechanical_reduction = motor_options["mechanicalReduction"] + encoder_to_joint_ratio.append(1.0 / mechanical_reduction) + except IndexError: + encoder_to_joint_ratio.append(1.0) + self.encoder_to_joint_ratio = np.array(encoder_to_joint_ratio) # Call `refresh_observation` manually to pre-compile it if necessary if not self._is_compiled: @@ -736,10 +762,14 @@ def fieldnames(self) -> List[List[str]]: for e in ("x", "y", "z", "w")] def refresh_observation(self, measurement: BaseObsT) -> None: - # Estimate the theoretical configuration of the robot from encoder data - self._q_th[self.encoder_to_config] = self.encoder_data + # Translate encoder data at joint level + joint_positions = self.encoder_to_joint_ratio * self.encoder_data + + # Update the configuration of the theoretical model of the robot + self._q_th[self.encoder_to_position_map] = joint_positions - # Update kinematic quantities according to the estimated configuration + # Update kinematic quantities according to the estimated configuration. + # FIXME: Compute frame placement only for relevant IMUs. pin.framesForwardKinematics( self.pinocchio_model_th, self.pinocchio_data_th, self._q_th) diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py index 606b80ced..5576a88a3 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py @@ -8,8 +8,7 @@ import numba as nb import gymnasium as gym -from jiminy_py.core import ( # pylint: disable=no-name-in-module - ImuSensor as imu) +from jiminy_py.core import ImuSensor # pylint: disable=no-name-in-module from ..bases import BaseObsT, BaseActT, BaseObserverBlock, InterfaceJiminyEnv from ..utils import (fill, @@ -200,7 +199,7 @@ def __init__(self, Optional: `1` by default. """ # Handling of default argument(s) - num_imu_sensors = len(env.robot.sensor_names[imu.type]) + num_imu_sensors = len(env.robot.sensors[ImuSensor.type]) if isinstance(kp, float): kp = np.full((num_imu_sensors,), kp) if isinstance(ki, float): @@ -268,7 +267,7 @@ def _initialize_state_space(self) -> None: # internal state of the (partially observable) MDP since the previous # observation must be provided anyway when integrating the observable # dynamics by definition. - num_imu_sensors = len(self.env.robot.sensor_names[imu.type]) + num_imu_sensors = len(self.env.robot.sensors[ImuSensor.type]) self.state_space = gym.spaces.Box( low=np.full((3, num_imu_sensors), -np.inf), high=np.full((3, num_imu_sensors), np.inf), @@ -281,7 +280,7 @@ def _initialize_observation_space(self) -> None: the robot at once, with special treatment for their twist part. See `__init__` documentation for details. """ - num_imu_sensors = len(self.env.robot.sensor_names[imu.type]) + num_imu_sensors = len(self.env.robot.sensors[ImuSensor.type]) self.observation_space = gym.spaces.Box( low=np.full((4, num_imu_sensors), -1.0 - 1e-9), high=np.full((4, num_imu_sensors), 1.0 + 1e-9), @@ -301,7 +300,7 @@ def _setup(self) -> None: # Refresh gyroscope and accelerometer proxies self.gyro, self.acc = np.split( - self.env.sensor_measurements[imu.type], 2) + self.env.sensor_measurements[ImuSensor.type], 2) # Reset the sensor bias fill(self._bias, 0) @@ -332,8 +331,8 @@ def get_state(self) -> np.ndarray: @property def fieldnames(self) -> List[List[str]]: - sensor_names = self.env.robot.sensor_names[imu.type] - return [[f"{name}.Quat{e}" for name in sensor_names] + imu_sensors = self.env.robot.sensors[ImuSensor.type] + return [[f"{sensor.name}.Quat{e}" for sensor in imu_sensors] for e in ("x", "y", "z", "w")] def refresh_observation(self, measurement: BaseObsT) -> None: @@ -358,9 +357,8 @@ def refresh_observation(self, measurement: BaseObsT) -> None: # Get true orientation of IMU frames imu_rots = [] robot = self.env.robot - for name in robot.sensor_names[imu.type]: - sensor = robot.get_sensor(imu.type, name) - assert isinstance(sensor, imu) + for sensor in robot.sensors[ImuSensor.type]: + assert isinstance(sensor, ImuSensor) rot = robot.pinocchio_data.oMf[sensor.frame_index].rotation imu_rots.append(rot) diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/motor_safety_limit.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/motor_safety_limit.py index 6c17a6b72..8654c82ff 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/motor_safety_limit.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/motor_safety_limit.py @@ -7,8 +7,7 @@ import numpy as np import numba as nb -from jiminy_py.core import ( # pylint: disable=no-name-in-module - EncoderSensor as encoder) +from jiminy_py.core import EncoderSensor # pylint: disable=no-name-in-module from .proportional_derivative_controller import get_encoder_to_motor_map @@ -113,18 +112,18 @@ def __init__(self, *, kp: float, kd: float, - soft_position_margin: float = 0.0, - soft_velocity_max: float = float("inf")) -> None: + soft_position_margin: float, + soft_velocity_max: float) -> None: """ :param name: Name of the block. :param env: Environment to connect with. :param kp: Scale of the velocity bound triggered by position limits. :param kd: Scale of the effort bound triggered by velocity limits. - :param soft_position_margin: Minimum distance of the current motor + :param soft_position_margin: Minimum distance of the current joint positions from their respective bounds before starting to break. - :param soft_velocity_max: Maximum velocity of the motor before + :param soft_velocity_max: Maximum velocity of the joint before starting to break. """ # Make sure that no other controller has been added prior to this block @@ -139,26 +138,33 @@ def __init__(self, self.kp = kp self.kd = kd + # Refresh mechanical reduction ratio + encoder_to_joint_ratio = [] + for motor in env.robot.motors: + motor_options = motor.get_options() + encoder_to_joint_ratio.append(motor_options["mechanicalReduction"]) + # Define buffers storing information about the motors for efficiency - motor_position_indices: List[int] = sum( - env.robot.motor_position_indices, []) - self.motors_position_lower = env.robot.position_limit_lower[ - motor_position_indices] + soft_position_margin - self.motors_position_upper = env.robot.position_limit_upper[ - motor_position_indices] - soft_position_margin - self.motors_velocity_limit = np.minimum(env.robot.velocity_limit[ - env.robot.motor_velocity_indices], soft_velocity_max) - self.motors_effort_limit = env.robot.command_limit[ - env.robot.motor_velocity_indices] + self.motors_position_lower = np.array([ + motor.position_limit_lower + ratio * soft_position_margin + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + self.motors_position_upper = np.array([ + motor.position_limit_upper - ratio * soft_position_margin + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + self.motors_velocity_limit = np.array([ + min(motor.velocity_limit, ratio * soft_velocity_max) + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + self.motors_effort_limit = np.array([ + motor.effort_limit for motor in env.robot.motors]) self.motors_effort_limit[ self.motors_position_lower > self.motors_position_upper] = 0.0 # Mapping from motors to encoders - self.encoder_to_motor = get_encoder_to_motor_map(env.robot) + self.encoder_to_motor_map = get_encoder_to_motor_map(env.robot) # Whether stored reference to encoder measurements are already in the # same order as the motors, allowing skipping re-ordering entirely. - self._is_same_order = isinstance(self.encoder_to_motor, slice) + self._is_same_order = isinstance(self.encoder_to_motor_map, slice) if not self._is_same_order: warnings.warn( "Consider using the same ordering for encoders and motors for " @@ -184,12 +190,12 @@ def _setup(self) -> None: # Refresh measured motor positions and velocities proxies self.q_measured, self.v_measured = ( - self.env.sensor_measurements[encoder.type]) + self.env.sensor_measurements[EncoderSensor.type]) @property def fieldnames(self) -> List[str]: - return [f"currentMotorTorque{name}" - for name in self.env.robot.motor_names] + return [f"currentMotorTorque{motor.name}" + for motor in self.env.robot.motors] def compute_command(self, action: np.ndarray, @@ -204,8 +210,8 @@ def compute_command(self, # Extract motor positions and velocity from encoder data q_measured, v_measured = self.q_measured, self.v_measured if not self._is_same_order: - q_measured = q_measured[self.encoder_to_motor] - v_measured = v_measured[self.encoder_to_motor] + q_measured = q_measured[self.encoder_to_motor_map] + v_measured = v_measured[self.encoder_to_motor_map] # Clip command according to safe effort bounds apply_safety_limits(action, diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py index 1a5204051..ce5fb810f 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py @@ -10,8 +10,7 @@ import jiminy_py.core as jiminy from jiminy_py.core import ( # pylint: disable=no-name-in-module - array_copyto, - EncoderSensor as encoder) + EncoderSensor, array_copyto) from ..bases import BaseObsT, InterfaceJiminyEnv, BaseControllerBlock from ..utils import fill @@ -235,27 +234,27 @@ def get_encoder_to_motor_map(robot: jiminy.Robot) -> Union[slice, List[int]]: :returns: A slice if possible, a list of indices otherwise. """ - # Define the mapping from motors to encoders - encoder_to_motor = [-1 for _ in range(robot.nmotors)] - encoders = [robot.get_sensor(encoder.type, sensor_name) - for sensor_name in robot.sensor_names[encoder.type]] - for i, motor_name in enumerate(robot.motor_names): - motor = robot.get_motor(motor_name) - for j, sensor in enumerate(encoders): - assert isinstance(sensor, encoder) - if motor.joint_index == sensor.joint_index: - encoder_to_motor[sensor.index] = i - encoders.pop(j) + # Loop over all actuated joints and look for their encoder counterpart + encoder_to_motor_map = [-1,] * robot.nmotors + encoders = robot.sensors[EncoderSensor.type] + for motor in robot.motors: + for i, sensor in enumerate(encoders): + assert isinstance(sensor, EncoderSensor) + if motor.index == sensor.motor_index: + encoder_to_motor_map[sensor.index] = motor.index + del encoders[i] break else: raise RuntimeError( - f"No encoder sensor associated with motor '{motor_name}'. " - "Every actuated joint must have encoder sensors attached.") + f"No encoder sensor associated with motor '{motor.name}'. " + "Every actuated joint must have a encoder sensor attached on " + "motor side.") # Try converting it to slice if possible - if (np.array(encoder_to_motor) == np.arange(robot.nmotors)).all(): + if (np.array(encoder_to_motor_map) == np.arange(robot.nmotors)).all(): return slice(None) - return encoder_to_motor + + return encoder_to_motor_map class PDController( @@ -288,9 +287,9 @@ def __init__(self, update_ratio: int = 1, kp: Union[float, List[float], np.ndarray], kd: Union[float, List[float], np.ndarray], - target_position_margin: float = 0.0, - target_velocity_limit: float = float("inf"), - target_acceleration_limit: float = float("inf")) -> None: + joint_position_margin: float = 0.0, + joint_velocity_limit: float = float("inf"), + joint_acceleration_limit: float = float("inf")) -> None: """ :param name: Name of the block. :param env: Environment to connect with. @@ -300,16 +299,14 @@ def __init__(self, Optional: 1 by default. :param kp: PD controller position-proportional gains in motor order. :param kd: PD controller velocity-proportional gains in motor order. - :param target_position_margin: Minimum distance of the motor target - positions from their respective bounds. - Optional: 0.0 by default. - :param target_velocity_limit: Restrict maximum motor target velocities - wrt their hardware specifications. - Optional: "inf" by default. - :param target_acceleration_limit: - Restrict maximum motor target accelerations wrt their hardware - specifications. - Optional: "inf" by default. + :param joint_position_margin: Minimum distance of the joint target + positions from their respective bounds. + Optional: 0.0 by default. + :param joint_velocity_limit: Restrict maximum joint target velocities + wrt their hardware specifications. + Optional: "inf" by default. + :param joint_acceleration_limit: Maximum joint target acceleration. + Optional: "inf" by default. """ # Make sure the action space of the environment has not been altered if env.action_space is not env.unwrapped.action_space: @@ -330,11 +327,11 @@ def __init__(self, self.kd = kd # Mapping from motors to encoders - self.encoder_to_motor = get_encoder_to_motor_map(env.robot) + self.encoder_to_motor_map = get_encoder_to_motor_map(env.robot) # Whether stored reference to encoder measurements are already in the # same order as the motors, allowing skipping re-ordering entirely. - self._is_same_order = isinstance(self.encoder_to_motor, slice) + self._is_same_order = isinstance(self.encoder_to_motor_map, slice) if not self._is_same_order: warnings.warn( "Consider using the same ordering for encoders and motors for " @@ -345,42 +342,43 @@ def __init__(self, # to another, the observation and action spaces are required to stay # the same the whole time. This induces that the motors effort limit # must not change unlike the mapping from full state to motors. - self.motors_effort_limit = env.robot.command_limit[ - env.robot.motor_velocity_indices] - - # Extract the motors target position bounds from the model - motors_position_lower: List[float] = [] - motors_position_upper: List[float] = [] - for motor_name in env.robot.motor_names: - motor = env.robot.get_motor(motor_name) - joint_type = jiminy.get_joint_type( - env.robot.pinocchio_model, motor.joint_index) - if joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: - lower, upper = float("-inf"), float("inf") - else: - motor_position_index = motor.joint_position_index - lower = env.robot.position_limit_lower[motor_position_index] - upper = env.robot.position_limit_upper[motor_position_index] - motors_position_lower.append(lower + target_position_margin) - motors_position_upper.append(upper - target_position_margin) - - # Extract the motors target velocity bounds - motors_velocity_limit = np.minimum( - env.robot.velocity_limit[env.robot.motor_velocity_indices], - target_velocity_limit) - - # Compute acceleration bounds allowing unrestricted bang-bang control + self.motors_effort_limit = np.array([ + motor.effort_limit for motor in env.robot.motors]) + + # Refresh mechanical reduction ratio + encoder_to_joint_ratio = [] + for motor in env.robot.motors: + motor_options = motor.get_options() + encoder_to_joint_ratio.append(motor_options["mechanicalReduction"]) + + # Define the motors target position bounds + motors_position_lower = np.array([ + motor.position_limit_lower + ratio * joint_position_margin + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + motors_position_upper = np.array([ + motor.position_limit_upper - ratio * joint_position_margin + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + + # Define the motors target velocity bounds + motors_velocity_limit = np.array([ + min(motor.velocity_limit, ratio * joint_velocity_limit) + for motor, ratio in zip(env.robot.motors, encoder_to_joint_ratio)]) + + # Define acceleration bounds allowing unrestricted bang-bang control range_limit = 2 * motors_velocity_limit / env.step_dt effort_limit = self.motors_effort_limit / ( self.kp * env.step_dt * np.maximum(env.step_dt / 2, self.kd)) + target_acceleration_limit = np.array([ + ratio * joint_acceleration_limit + for ratio in encoder_to_joint_ratio]) acceleration_limit = np.minimum( np.minimum(range_limit, effort_limit), target_acceleration_limit) # Compute command state bounds - self._command_state_lower = np.stack([np.array(motors_position_lower), + self._command_state_lower = np.stack([motors_position_lower, -motors_velocity_limit, -acceleration_limit], axis=0) - self._command_state_upper = np.stack([np.array(motors_position_upper), + self._command_state_upper = np.stack([motors_position_upper, motors_velocity_limit, acceleration_limit], axis=0) @@ -429,15 +427,15 @@ def _setup(self) -> None: # Refresh measured motor positions and velocities proxies self.q_measured, self.v_measured = ( - self.env.sensor_measurements[encoder.type]) + self.env.sensor_measurements[EncoderSensor.type]) # Reset the command state fill(self._command_state, 0) @property def fieldnames(self) -> List[str]: - return [f"currentTarget{N_ORDER_DERIVATIVE_NAMES[2]}{name}" - for name in self.env.robot.motor_names] + return [f"currentTarget{N_ORDER_DERIVATIVE_NAMES[2]}{motor.name}" + for motor in self.env.robot.motors] def get_state(self) -> np.ndarray: return self._command_state[:2] @@ -470,8 +468,8 @@ def compute_command(self, action: np.ndarray, command: np.ndarray) -> None: # Extract motor positions and velocity from encoder data q_measured, v_measured = self.q_measured, self.v_measured if not self._is_same_order: - q_measured = q_measured[self.encoder_to_motor] - v_measured = v_measured[self.encoder_to_motor] + q_measured = q_measured[self.encoder_to_motor_map] + v_measured = v_measured[self.encoder_to_motor_map] # Update target motor accelerations array_copyto(self._command_acceleration, action) @@ -547,9 +545,6 @@ def __init__(self, # Define some proxies for convenience self._pd_controller = controller - # Allocate memory for the target motor accelerations - self._target_accelerations = np.zeros((env.robot.nmotors,)) - # Initialize the controller super().__init__(name, env, update_ratio) @@ -564,17 +559,10 @@ def _initialize_action_space(self) -> None: high=self._pd_controller._command_state_upper[self.order], dtype=np.float64) - def _setup(self) -> None: - # Call base implementation - super()._setup() - - # Reset the target motor accelerations - fill(self._target_accelerations, 0) - @property def fieldnames(self) -> List[str]: - return [f"nextTarget{N_ORDER_DERIVATIVE_NAMES[self.order]}{name}" - for name in self.env.robot.motor_names] + return [f"nextTarget{N_ORDER_DERIVATIVE_NAMES[self.order]}{motor.name}" + for motor in self.env.robot.motors] def compute_command(self, action: np.ndarray, command: np.ndarray) -> None: """Compute the target motor accelerations from the desired value of diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py new file mode 100644 index 000000000..2558f282a --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py @@ -0,0 +1,15 @@ +# pylint: disable=missing-module-docstring + +from .generic import (radial_basis_function, + AdditiveMixtureReward, + MultiplicativeMixtureReward, + SurviveReward) +from .locomotion import OdometryVelocityReward + +__all__ = [ + "radial_basis_function", + "AdditiveMixtureReward", + "MultiplicativeMixtureReward", + "SurviveReward", + "OdometryVelocityReward" +] diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py new file mode 100644 index 000000000..1a1c487de --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py @@ -0,0 +1,218 @@ +"""Generic reward components that may be relevant for any kind of robot, +regardless its topology (multiple or single branch, fixed or floating base...) +and the application (locomotion, grasping...). +""" +import math +import logging +from typing import Sequence, Optional + +import numpy as np +import numba as nb + +from ..bases import ( + InfoType, InterfaceJiminyEnv, AbstractReward, BaseMixtureReward) + + +# Reward value at cutoff threshold +RBF_CUTOFF_ESP = 1.0e-2 + + +LOGGER = logging.getLogger(__name__) + + +@nb.jit(nopython=True, cache=True, fastmath=True) +def radial_basis_function(error: float, + cutoff: float, + order: int = 2) -> float: + r"""Radial basis function (RBF) kernel (aka squared-exponential kernel). + + The RBF kernel is defined as: + + .. math:: + + f(x) = \exp{\frac{dist(x, x_ref)}{2 \sigma^2}} + + where :math:`dist(x, x_ref)` is some distance metric of the error between + the observed (:math:`x`) and desired (:math:`x_ref`) values of a + multi-variate quantity. The L2-norm (Euclidean norm) was used when it was + first introduced as a non-linear kernel for Support Vector Machine (SVM) + algorithm. Such restriction does not make sense in the context of reward + normalization. The scaling parameter :math:`sigma` is derived from the + user-specified cutoff. The latter is defined as the distance from which the + attenuation reaches 99%. + + :param error: Multi-variate error on some tangent space. + :param cutoff: Cut-off threshold to consider. + :param order: Order of Lp-Norm that will be used as distance metric. + """ + distance = np.linalg.norm(error, order) + return math.pow(RBF_CUTOFF_ESP, math.pow(distance / cutoff, 2)) + + +class AdditiveMixtureReward(BaseMixtureReward): + """Weighted sum of multiple independent reward components. + + Aggregation of reward components using the addition operator is suitable + when improving the behavior for any of them without the others is equally + beneficial, and unbalanced performance for each reward component is + considered acceptable rather than detrimental. It especially makes sense + for reward that are not competing with each other (improving one tends to + impede some other). In the latter case, the multiplicative operator is more + appropriate. See `MultiplicativeMixtureReward` documentation for details. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + name: str, + components: Sequence[AbstractReward], + weights: Optional[Sequence[float]] = None) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param name: Desired name of the total reward. + :param components: Sequence of reward components to aggregate. + :param weights: Sequence of weights associated with each reward + components, with the same ordering as 'components'. + Optional: 1.0 for all reward components by default. + """ + # Handling of default arguments + if weights is None: + weights = (1.0,) * len(components) + + # Make sure that the weight sequence is consistent with the components + if len(weights) != len(components): + raise ValueError( + "Exactly one weight per reward component must be specified.") + + # Determine whether the cumulative reward is normalized + weight_total = 0.0 + for weight, reward in zip(weights, components): + if not reward.is_normalized: + LOGGER.warning( + "Reward '%s' is not normalized. Aggregating rewards that " + "are not normalized using the addition operator is not " + "recommended.", reward.name) + is_normalized = False + break + weight_total += weight + else: + is_normalized = abs(weight_total - 1.0) < 1e-4 + + # Backup user-arguments + self.weights = weights + + # Call base implementation + super().__init__(env, name, components, self._reduce, is_normalized) + + def _reduce(self, values: Sequence[Optional[float]]) -> Optional[float]: + """Compute the weighted sum of all the reward components that has been + evaluated, filtering out the others. + + This method returns `None` if no reward component has been evaluated. + + :param values: Sequence of scalar value for reward components that has + been evaluated, `None` otherwise, with the same ordering + as 'components'. + + :returns: Scalar value if at least one of the reward component has been + evaluated, `None` otherwise. + """ + # TODO: x2 speedup can be expected with `nb.jit` + total, any_value = 0.0, False + for weight, value in zip(self.weights, values): + if value is not None: + total += weight * value + any_value = True + return total if any_value else None + + +AdditiveMixtureReward.is_normalized.__doc__ = \ + """Whether the reward is guaranteed to be normalized, ie it is in range + [0.0, 1.0]. + + The cumulative reward is considered normalized if all its individual + reward components are normalized and their weights sums up to 1.0. + """ + + +class MultiplicativeMixtureReward(BaseMixtureReward): + """Product of multiple independent reward components. + + Aggregation of reward components using multiplication operator is suitable + when maintaining balanced performance between all reward components is + essential, and having poor performance for any of them is unacceptable. + This type of aggregation is especially useful when reward components are + competing with each other (improving one tends to impede some other) but + not mutually exclusive. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + name: str, + components: Sequence[AbstractReward] + ) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param name: Desired name of the reward. + :param components: Sequence of reward components to aggregate. + """ + # Determine whether the cumulative reward is normalized + is_normalized = all(reward.is_normalized for reward in components) + + # Call base implementation + super().__init__(env, name, components, self._reduce, is_normalized) + + def _reduce(self, values: Sequence[Optional[float]]) -> Optional[float]: + """Compute the product of all the reward components that has been + evaluated, filtering out the others. + + This method returns `None` if no reward component has been evaluated. + + :param values: Sequence of scalar value for reward components that has + been evaluated, `None` otherwise, with the same ordering + as 'components'. + + :returns: Scalar value if at least one of the reward component has been + evaluated, `None` otherwise. + """ + # TODO: x2 speedup can be expected with `nb.jit` + total, any_value = 1.0, False + for value in values: + if value is not None: + total *= value + any_value = True + return total if any_value else None + + +AdditiveMixtureReward.is_normalized.__doc__ = \ + """Whether the reward is guaranteed to be normalized, ie it is in range + [0.0, 1.0]. + + The cumulative reward is considered normalized if all its individual + reward components are normalized. + """ + + +class SurviveReward(AbstractReward): + """Constant positive reward equal to 1.0 systematically, unless the current + state of the environment is the terminal state. In which case, the value + 0.0 is returned instead. + """ + + def __init__(self, env: InterfaceJiminyEnv) -> None: + """ + :param env: Base or wrapped jiminy environment. + """ + super().__init__(env, "reward_survive") + + @property + def is_terminal(self) -> Optional[bool]: + return False + + @property + def is_normalized(self) -> bool: + return True + + def compute(self, terminated: bool, info: InfoType) -> Optional[float]: + """Return a constant positive reward equal to 1.0 no matter what. + """ + return 1.0 diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py new file mode 100644 index 000000000..997208425 --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py @@ -0,0 +1,44 @@ +"""Rewards mainly relevant for locomotion tasks on floating-base robots. +""" +from operator import sub +from functools import partial + +from ..bases import InterfaceJiminyEnv, BaseQuantityReward, QuantityEvalMode +from ..quantities import AverageOdometryVelocity, BinaryOpQuantity + +from .generic import radial_basis_function + + +class OdometryVelocityReward(BaseQuantityReward): + """Reward the agent for tracking a reference odometry velocity. + + A reference trajectory must be selected before evaluating this reward + otherwise an exception will be risen. See `DatasetTrajectoryQuantity` and + `AbstractQuantity` documentations for details. + + The error transform in a normalized reward to maximize by applying RBF + kernel on the error. The reward will be 0.0 if the error cancels out + completely and less than 0.01 above the user-specified cutoff threshold. + """ + def __init__(self, + env: InterfaceJiminyEnv, + cutoff: float) -> None: + """ + :param cutoff: Cutoff threshold for the RBF kernel transform. + """ + # Backup some user argument(s) + self.cutoff = cutoff + + # Call base implementation + super().__init__( + env, + "reward_odometry_velocity", + (BinaryOpQuantity, dict( + quantity_left=(AverageOdometryVelocity, dict( + mode=QuantityEvalMode.TRUE)), + quantity_right=(AverageOdometryVelocity, dict( + mode=QuantityEvalMode.REFERENCE)), + op=sub)), + partial(radial_basis_function, cutoff=self.cutoff, order=2), + is_normalized=True, + is_terminal=False) diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py index 569e987e0..656d2925e 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py @@ -23,12 +23,7 @@ import jiminy_py.core as jiminy from jiminy_py import tree from jiminy_py.core import ( # pylint: disable=no-name-in-module - array_copyto, - EncoderSensor as encoder, - EffortSensor as effort, - ContactSensor as contact, - ForceSensor as force, - ImuSensor as imu) + EncoderSensor, EffortSensor, array_copyto) from jiminy_py.dynamics import compute_freeflyer_state_from_fixed_body from jiminy_py.log import extract_variables_from_log from jiminy_py.simulator import Simulator, TabbedFigure @@ -36,7 +31,7 @@ interactive_mode, get_default_backend, Viewer) -from jiminy_py.viewer.replay import viewer_lock # type: ignore[attr-defined] +from jiminy_py.viewer.replay import viewer_lock import pinocchio as pin @@ -49,7 +44,8 @@ build_contains, get_fieldnames, register_variables) -from ..bases import (ObsT, +from ..bases import (DT_EPS, + ObsT, ActT, InfoType, SensorMeasurementStackMap, @@ -61,24 +57,11 @@ # Maximum realtime slowdown of simulation steps before triggering timeout error -TIMEOUT_RATIO = 10 +TIMEOUT_RATIO = 15 # Absolute tolerance when checking that observations are valid OBS_CONTAINS_TOL = 0.01 -# Define universal bounds for the observation space -FREEFLYER_POS_TRANS_MAX = 1000.0 -FREEFLYER_VEL_LIN_MAX = 1000.0 -FREEFLYER_VEL_ANG_MAX = 10000.0 -JOINT_POS_MAX = 10000.0 -JOINT_VEL_MAX = 100.0 -FLEX_VEL_ANG_MAX = 10000.0 -MOTOR_EFFORT_MAX = 1000.0 -SENSOR_FORCE_MAX = 100000.0 -SENSOR_MOMENT_MAX = 10000.0 -SENSOR_GYRO_MAX = 100.0 -SENSOR_ACCEL_MAX = 10000.0 - LOGGER = logging.getLogger(__name__) @@ -119,7 +102,6 @@ class BaseJiminyEnv(InterfaceJiminyEnv[ObsT, ActT], def __init__(self, simulator: Simulator, step_dt: float, - enforce_bounded_spaces: bool = False, debug: bool = False, render_mode: Optional[str] = None, **kwargs: Any) -> None: @@ -135,11 +117,6 @@ def __init__(self, snapshot as an RGB array without showing it on the screen. Optional: 'human' by default if available with the current backend (or default if none), 'rgb_array' otherwise. - :param enforce_bounded_spaces: - Whether to enforce finite bounds for the observation and action - spaces. If so, then '\*_MAX' are used whenever it is necessary. - Note that whose bounds are very spread to make sure it is suitable - for the vast majority of systems. :param debug: Whether the debug mode must be enabled. Doing it enables telemetry recording. :param render_mode: Desired rendering mode, ie "human" or "rgb_array". @@ -193,7 +170,6 @@ def __init__(self, self.simulator: Simulator = simulator self._step_dt = step_dt self.render_mode = render_mode - self.enforce_bounded_spaces = enforce_bounded_spaces self.debug = debug # Define some proxies for fast access. @@ -235,8 +211,7 @@ def __init__(self, self.total_reward = 0.0 # Number of simulation steps performed - self.num_steps = -1 - self.max_steps = 0 + self.num_steps = np.array(-1, dtype=np.int64) self._num_steps_beyond_terminate: Optional[int] = None # Initialize the interfaces through multiple inheritance @@ -298,7 +273,8 @@ def __init__(self, action_size = self.action.size if action_size > 0 and action_size == self.robot.nmotors: action_fieldnames = [ - ".".join(('action', e)) for e in self.robot.motor_names] + ".".join(('action', motor.name)) + for motor in self.robot.motors] self.register_variable( 'action', self.action, action_fieldnames) @@ -336,38 +312,28 @@ def _get_time_space(self) -> spaces.Box: shape=(), dtype=np.float64) - def _get_agent_state_space( - self, use_theoretical_model: bool = False) -> spaces.Dict: + def _get_agent_state_space(self, + use_theoretical_model: bool = False, + ignore_velocity_limit: bool = True + ) -> spaces.Dict: """Get state space. This method is not meant to be overloaded in general since the definition of the state space is mostly consensual. One must rather overload `_initialize_observation_space` to customize the observation space as a whole. + + :param use_theoretical_model: Whether to compute the state space + associated with the theoretical model + instead of the extended simulation model. + :param ignore_velocity_limit: Whether to ignore the velocity bounds + specified in model. """ # Define some proxies for convenience - model_options = self.robot.get_model_options() - joint_velocity_indices = self.robot.mechanical_joint_velocity_indices - position_limit_upper = self.robot.position_limit_upper - position_limit_lower = self.robot.position_limit_lower - velocity_limit = self.robot.velocity_limit - - # Replace inf bounds of the state space if requested - if self.enforce_bounded_spaces: - if self.robot.has_freeflyer: - position_limit_lower[:3] = -FREEFLYER_POS_TRANS_MAX - position_limit_upper[:3] = +FREEFLYER_POS_TRANS_MAX - velocity_limit[:3] = FREEFLYER_VEL_LIN_MAX - velocity_limit[3:6] = FREEFLYER_VEL_ANG_MAX - - for joint_index in self.robot.flexibility_joint_indices: - joint_velocity_index = ( - self.robot.pinocchio_model.joints[joint_index].idx_v) - velocity_limit[ - joint_velocity_index + np.arange(3)] = FLEX_VEL_ANG_MAX - - if not model_options['joints']['enableVelocityLimit']: - velocity_limit[joint_velocity_indices] = JOINT_VEL_MAX + pinocchio_model = self.robot.pinocchio_model + position_limit_lower = pinocchio_model.lowerPositionLimit + position_limit_upper = pinocchio_model.upperPositionLimit + velocity_limit = pinocchio_model.velocityLimit # Deduce bounds associated the theoretical model from the extended one if use_theoretical_model: @@ -377,13 +343,18 @@ def _get_agent_state_space( velocity_limit = self.robot.get_theoretical_velocity_from_extended( velocity_limit) + # Ignore velocity bounds in requested + if ignore_velocity_limit: + velocity_limit = np.full_like(velocity_limit, float("inf")) + # Aggregate position and velocity bounds to define state space return spaces.Dict(OrderedDict( q=spaces.Box(low=position_limit_lower, high=position_limit_upper, dtype=np.float64), - v=spaces.Box(low=-velocity_limit, - high=velocity_limit, + v=spaces.Box(low=float("-inf"), + high=float("inf"), + shape=(self.robot.pinocchio_model.nv,), dtype=np.float64))) def _get_measurements_space(self) -> spaces.Dict: @@ -406,8 +377,7 @@ def _get_measurements_space(self) -> spaces.Dict: .. code-block:: python - sensor_name = env.robot.sensor_names[key][j] - sensor = env.robot.get_sensor(key, sensor_name) + sensor = env.robot.sensors[key][j] .. warning: This method is not meant to be overloaded in general since the @@ -416,20 +386,11 @@ def _get_measurements_space(self) -> spaces.Dict: observation space as a whole. """ # Define some proxies for convenience - sensor_measurements = self.robot.sensor_measurements - command_limit = self.robot.command_limit - position_space, velocity_space = self._get_agent_state_space().values() - assert isinstance(position_space, spaces.Box) - assert isinstance(velocity_space, spaces.Box) - - # Replace inf bounds of the action space - for motor_name in self.robot.motor_names: - motor = self.robot.get_motor(motor_name) - motor_options = motor.get_options() - if not motor_options["enableCommandLimit"]: - command_limit[motor.joint_velocity_index] = MOTOR_EFFORT_MAX + position_limit_lower = self.robot.pinocchio_model.lowerPositionLimit + position_limit_upper = self.robot.pinocchio_model.upperPositionLimit # Initialize the bounds of the sensor space + sensor_measurements = self.robot.sensor_measurements sensor_space_lower = OrderedDict( (key, np.full(value.shape, -np.inf)) for key, value in sensor_measurements.items()) @@ -438,72 +399,40 @@ def _get_measurements_space(self) -> spaces.Dict: for key, value in sensor_measurements.items()) # Replace inf bounds of the encoder sensor space - if encoder.type in sensor_measurements.keys(): - sensor_list = self.robot.sensor_names[encoder.type] - for sensor_name in sensor_list: - # Get the position and velocity bounds of the sensor. - # Note that for rotary unbounded encoders, the sensor bounds - # cannot be extracted from the configuration vector limits - # since the representation is different: cos/sin for the - # configuration, and principal value of the angle for the - # sensor. - sensor = self.robot.get_sensor(encoder.type, sensor_name) - assert isinstance(sensor, encoder) - sensor_index = sensor.index - joint = self.robot.pinocchio_model.joints[sensor.joint_index] - if sensor.joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: - sensor_position_lower = -np.pi - sensor_position_upper = np.pi - else: - sensor_position_lower = position_space.low[joint.idx_q] - sensor_position_upper = position_space.high[joint.idx_q] - sensor_velocity_limit = velocity_space.high[joint.idx_v] - - # Update the bounds accordingly - sensor_space_lower[encoder.type][:, sensor_index] = ( - sensor_position_lower, -sensor_velocity_limit) - sensor_space_upper[encoder.type][:, sensor_index] = ( - sensor_position_upper, sensor_velocity_limit) + for sensor in self.robot.sensors.get(EncoderSensor.type, ()): + # Get the position bounds of the sensor. + # Note that for rotary unbounded encoders, the sensor bounds + # cannot be extracted from the motor because only the principal + # value of the angle is observed by the sensor. + assert isinstance(sensor, EncoderSensor) + joint = self.robot.pinocchio_model.joints[sensor.joint_index] + joint_type = jiminy.get_joint_type(joint) + if joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: + sensor_position_lower = - np.pi + sensor_position_upper = np.pi + else: + try: + motor = self.robot.motors[sensor.motor_index] + sensor_position_lower = motor.position_limit_lower + sensor_position_upper = motor.position_limit_upper + except IndexError: + sensor_position_lower = position_limit_lower[joint.idx_q] + sensor_position_upper = position_limit_upper[joint.idx_q] + + # Update the bounds accordingly + sensor_space_lower[EncoderSensor.type][0, sensor.index] = ( + sensor_position_lower) + sensor_space_upper[EncoderSensor.type][0, sensor.index] = ( + sensor_position_upper) # Replace inf bounds of the effort sensor space - if effort.type in sensor_measurements.keys(): - sensor_list = self.robot.sensor_names[effort.type] - for sensor_name in sensor_list: - sensor = self.robot.get_sensor(effort.type, sensor_name) - assert isinstance(sensor, effort) - sensor_index = sensor.index - motor_velocity_index = self.robot.motor_velocity_indices[ - sensor.motor_index] - sensor_space_lower[effort.type][0, sensor_index] = ( - -command_limit[motor_velocity_index]) - sensor_space_upper[effort.type][0, sensor_index] = ( - command_limit[motor_velocity_index]) - - # Replace inf bounds of the imu sensor space - if self.enforce_bounded_spaces: - # Replace inf bounds of the contact sensor space - if contact.type in sensor_measurements.keys(): - sensor_space_lower[contact.type][:] = -SENSOR_FORCE_MAX - sensor_space_upper[contact.type][:] = SENSOR_FORCE_MAX - - # Replace inf bounds of the force sensor space - if force.type in sensor_measurements.keys(): - sensor_space_lower[force.type][:3] = -SENSOR_FORCE_MAX - sensor_space_upper[force.type][:3] = SENSOR_FORCE_MAX - sensor_space_lower[force.type][3:] = -SENSOR_MOMENT_MAX - sensor_space_upper[force.type][3:] = SENSOR_MOMENT_MAX - - # Replace inf bounds of the imu sensor space - if imu.type in sensor_measurements.keys(): - gyro_index = [ - field.startswith('Gyro') for field in imu.fieldnames] - sensor_space_lower[imu.type][gyro_index] = -SENSOR_GYRO_MAX - sensor_space_upper[imu.type][gyro_index] = SENSOR_GYRO_MAX - - accel_index = [ - field.startswith('Accel') for field in imu.fieldnames] - sensor_space_lower[imu.type][accel_index] = -SENSOR_ACCEL_MAX - sensor_space_upper[imu.type][accel_index] = SENSOR_ACCEL_MAX + for sensor in self.robot.sensors.get(EffortSensor.type, ()): + assert isinstance(sensor, EffortSensor) + motor = self.robot.motors[sensor.motor_index] + sensor_space_lower[EffortSensor.type][0, sensor.index] = ( + - motor.effort_limit) + sensor_space_upper[EffortSensor.type][0, sensor.index] = ( + motor.effort_limit) return spaces.Dict(OrderedDict( (key, spaces.Box(low=min_val, high=max_val, dtype=np.float64)) @@ -522,21 +451,12 @@ def _initialize_action_space(self) -> None: robot is uniquely defined. """ # Get effort limit - command_limit = self.robot.command_limit - - # Replace inf bounds of the effort limit if requested - if self.enforce_bounded_spaces: - for motor_name in self.robot.motor_names: - motor = self.robot.get_motor(motor_name) - motor_options = motor.get_options() - if not motor_options["enableCommandLimit"]: - command_limit[motor.joint_velocity_index] = \ - MOTOR_EFFORT_MAX + command_limit = np.array([ + motor.effort_limit for motor in self.robot.motors]) # Set the action space - action_scale = command_limit[self.robot.motor_velocity_indices] self.action_space = spaces.Box( - low=-action_scale, high=action_scale, dtype=np.float64) + low=-command_limit, high=command_limit, dtype=np.float64) def _initialize_seed(self, seed: Optional[int] = None) -> None: """Specify the seed of the environment. @@ -716,7 +636,7 @@ def reset(self, # type: ignore[override] self.simulator.reset(remove_all_forces=False) # Reset some internal buffers - self.num_steps = 0 + self.num_steps[()] = 0 self._num_steps_beyond_terminate = None # Create a new log file @@ -798,9 +718,6 @@ def reset(self, # type: ignore[override] self.robot.controller = jiminy.FunctionalController( partial(type(env)._controller_handle, weakref.proxy(env))) - # Configure the maximum number of steps - self.max_steps = int(self.simulation_duration_max // self.step_dt) - # Register user-specified variables to the telemetry for header, value in self._registered_variables.values(): register_variables(self.robot.controller, header, value) @@ -848,8 +765,11 @@ def reset(self, # type: ignore[override] f"'nan' value found in observation ({obs}). Something " "went wrong with `refresh_observation` method.") + # Reset the extra information buffer + self._info.clear() + # The simulation cannot be done before doing a single step. - if any(self.has_terminated()): + if any(self.has_terminated(self._info)): raise RuntimeError( "The simulation has already terminated at `reset`. Check the " "implementation of `has_terminated` if overloaded.") @@ -933,6 +853,15 @@ def step(self, # type: ignore[override] self.simulator.stop() raise + # Make sure there is no 'nan' value in observation + if is_nan(self._robot_state_a): + raise RuntimeError( + "The acceleration of the system is 'nan'. Something went " + "wrong with jiminy engine.") + + # Update number of (successful) steps + self.num_steps += 1 + # Update shared buffers self._refresh_buffers() @@ -945,22 +874,16 @@ def step(self, # type: ignore[override] self._robot_state_v, self.robot.sensor_measurements) - # Make sure there is no 'nan' value in observation - if is_nan(self._robot_state_a): - raise RuntimeError( - "The acceleration of the system is 'nan'. Something went " - "wrong with jiminy engine.") - # Reset the extra information buffer self._info.clear() # Check if the simulation is over. # Note that 'truncated' is forced to True if the integration failed or # if the maximum number of steps will be exceeded next step. - terminated, truncated = self.has_terminated() + terminated, truncated = self.has_terminated(self._info) truncated = ( truncated or not self.is_simulation_running or - self.num_steps >= self.max_steps) + self.stepper_state.t + DT_EPS > self.simulation_duration_max) # Check if stepping after done and if it is an undefined behavior if self._num_steps_beyond_terminate is None: @@ -980,7 +903,7 @@ def step(self, # type: ignore[override] reward = float('nan') else: # Compute reward and update extra information - reward = self.compute_reward(terminated, truncated, self._info) + reward = self.compute_reward(terminated, self._info) # Make sure the reward is not 'nan' if math.isnan(reward): @@ -995,9 +918,6 @@ def step(self, # type: ignore[override] if self.debug and self._num_steps_beyond_terminate == 0: self.simulator.write_log(self.log_path, format="binary") - # Update number of (successful) steps - self.num_steps += 1 - # Clip (and copy) the most derived observation before returning it obs = self._get_clipped_env_observation() @@ -1099,7 +1019,7 @@ def replay(self, **kwargs: Any) -> None: kwargs['close_backend'] = not self.simulator.is_viewer_available # Stop any running simulation before replay if `has_terminated` is True - if self.is_simulation_running and any(self.has_terminated()): + if self.is_simulation_running and any(self.has_terminated({})): self.simulator.stop() with viewer_lock: @@ -1350,9 +1270,19 @@ def _setup(self) -> None: if self.debug: engine_options["stepper"]["timeout"] = 0.0 - # Force disabling logging of geometries unless in debug or eval modes - if self.is_training and not self.debug: - engine_options["telemetry"]["isPersistent"] = False + # Enable full logging in debug and evaluation mode + if self.debug or not self.is_training: + # Enable telemetry at engine-level + telemetry_options = engine_options["telemetry"] + for key in telemetry_options.keys(): + telemetry_options[key] = True + + # Enable telemetry at robot-level + robot_options = self.robot.get_options() + robot_telemetry_options = robot_options["telemetry"] + for key in robot_telemetry_options.keys(): + robot_telemetry_options[key] = True + self.robot.set_options(robot_options) # Update engine options self.simulator.set_options(engine_options) @@ -1396,8 +1326,8 @@ def _neutral(self) -> np.ndarray: q = pin.neutral(self.robot.pinocchio_model) # Make sure it is not out-of-bounds before returning - position_limit_lower = self.robot.position_limit_lower - position_limit_upper = self.robot.position_limit_upper + position_limit_lower = self.robot.pinocchio_model.lowerPositionLimit + position_limit_upper = self.robot.pinocchio_model.upperPositionLimit for idx, val in enumerate(q): lo, hi = position_limit_lower[idx], position_limit_upper[idx] if hi < val or val < lo: @@ -1422,8 +1352,8 @@ def _sample_state(self) -> Tuple[np.ndarray, np.ndarray]: q = self._neutral() # Make sure the configuration is not out-of-bound - q.clip(self.robot.position_limit_lower, - self.robot.position_limit_upper, + q.clip(self.robot.pinocchio_model.lowerPositionLimit, + self.robot.pinocchio_model.upperPositionLimit, out=q) # Make sure the configuration is normalized @@ -1548,7 +1478,9 @@ def compute_command(self, action: ActT, command: np.ndarray) -> None: assert isinstance(action, np.ndarray) array_copyto(command, action) - def has_terminated(self) -> Tuple[bool, bool]: + def has_terminated(self, + info: InfoType # pylint: disable=unused-argument + ) -> Tuple[bool, bool]: """Determine whether the episode is over, because a terminal state of the underlying MDP has been reached or an aborting condition outside the scope of the MDP has been triggered. @@ -1567,6 +1499,8 @@ def has_terminated(self) -> Tuple[bool, bool]: This method is called after `refresh_observation`, so that the internal buffer 'observation' is up-to-date. + :param info: Dictionary of extra information for monitoring. + :returns: terminated and truncated flags. """ # Make sure that a simulation is running diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/envs/locomotion.py index 2af385360..661a854d9 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/locomotion.py @@ -3,16 +3,16 @@ """ import os import pathlib -from typing import Optional, Dict, Union, Any, Type, Sequence, Tuple +from typing import Any, Dict, Optional, Tuple import numpy as np from jiminy_py.core import ( # pylint: disable=no-name-in-module - EncoderSensor as encoder, - EffortSensor as effort, - ContactSensor as contact, - ForceSensor as force, - ImuSensor as imu, + EncoderSensor, + EffortSensor, + ContactSensor, + ForceSensor, + ImuSensor, PeriodicGaussianProcess, Robot) from jiminy_py.robot import BaseJiminyRobot @@ -38,25 +38,25 @@ FLEX_DAMPING_SCALE = 10 SENSOR_DELAY_SCALE = { - encoder.type: 3.0e-3, - effort.type: 0.0, - contact.type: 0.0, - force.type: 0.0, - imu.type: 0.0 + EncoderSensor: 3.0e-3, + EffortSensor: 0.0, + ContactSensor: 0.0, + ForceSensor: 0.0, + ImuSensor: 0.0 } SENSOR_NOISE_SCALE = { - encoder.type: np.array([0.0, 0.02]), - effort.type: np.array([10.0]), - contact.type: np.array([2.0, 2.0, 2.0, 10.0, 10.0, 10.0]), - force.type: np.array([2.0, 2.0, 2.0]), - imu.type: np.array([0.0, 0.0, 0.0, 0.01, 0.01, 0.01, 0.2, 0.2, 0.2]) + EncoderSensor: np.array([0.0, 0.02]), + EffortSensor: np.array([10.0]), + ContactSensor: np.array([2.0, 2.0, 2.0, 10.0, 10.0, 10.0]), + ForceSensor: np.array([2.0, 2.0, 2.0]), + ImuSensor: np.array([0.0, 0.0, 0.0, 0.01, 0.01, 0.01, 0.2, 0.2, 0.2]) } SENSOR_BIAS_SCALE = { - encoder.type: np.array([0.0, 0.0]), - effort.type: np.array([0.0]), - contact.type: np.array([4.0, 4.0, 4.0, 20.0, 20.0, 20.0]), - force.type: np.array([4.0, 4.0, 4.0]), - imu.type: np.array([0.01, 0.01, 0.01, 0.02, 0.02, 0.02, 0.0, 0.0, 0.0]) + EncoderSensor: np.array([0.0, 0.0]), + EffortSensor: np.array([0.0]), + ContactSensor: np.array([4.0, 4.0, 4.0, 20.0, 20.0, 20.0]), + ForceSensor: np.array([4.0, 4.0, 4.0]), + ImuSensor: np.array([0.01, 0.01, 0.01, 0.02, 0.02, 0.02, 0.0, 0.0, 0.0]) } DEFAULT_SIMULATION_DURATION = 30.0 # (s) Default simulation duration @@ -80,7 +80,6 @@ def __init__(self, mesh_path_dir: Optional[str] = None, simulation_duration_max: float = DEFAULT_SIMULATION_DURATION, step_dt: float = DEFAULT_STEP_DT, - enforce_bounded_spaces: bool = False, reward_mixture: Optional[dict] = None, std_ratio: Optional[dict] = None, config_path: Optional[str] = None, @@ -101,9 +100,6 @@ def __init__(self, :param simulation_duration_max: Maximum duration of a simulation before returning done. :param step_dt: Simulation timestep for learning. - :param enforce_bounded_spaces: - Whether to enforce finite bounds for the observation and action - spaces. If so, '\*_MAX' are used whenever it is necessary. :param reward_mixture: Weighting factors of selected contributions to total reward. :param std_ratio: Relative standard deviation of selected contributions @@ -206,8 +202,7 @@ def __init__(self, simulator.import_options(config_path) # Initialize base class - super().__init__( - simulator, step_dt, enforce_bounded_spaces, debug, **kwargs) + super().__init__(simulator, step_dt, debug, **kwargs) def _setup(self) -> None: """Configure the environment. @@ -233,12 +228,10 @@ def _setup(self) -> None: "`WalkerJiminyEnv` only supports robots with freeflyer.") # Update some internal buffers used for computing the reward - motor_effort_limit = self.robot.pinocchio_model.effortLimit[ - self.robot.motor_velocity_indices] - motor_velocity_limit = self.robot.velocity_limit[ - self.robot.motor_velocity_indices] - self._power_consumption_max = sum( - motor_effort_limit * motor_velocity_limit) + self._power_consumption_max = 0.0 + for motor in self.robot.motors: + motor_power_max = motor.velocity_limit * motor.effort_limit + self._power_consumption_max += motor_power_max # Compute the height of the freeflyer in neutral configuration # TODO: Take into account the ground profile. @@ -253,8 +246,8 @@ def _setup(self) -> None: # computation and log replay. engine_options['telemetry']['enableConfiguration'] = True engine_options['telemetry']['enableVelocity'] = True - engine_options['telemetry']['enableForceExternal'] = \ - 'disturbance' in self.std_ratio.keys() + if 'disturbance' in self.std_ratio.keys(): + engine_options['telemetry']['enableForceExternal'] = True # ============= Add some stochasticity to the environment ============= @@ -267,25 +260,27 @@ def _setup(self) -> None: # Add sensor noise, bias and delay if 'sensors' in self.std_ratio.keys(): - sensor_classes: Sequence[Union[ - Type[encoder], Type[effort], Type[contact], Type[force], - Type[imu]]] = (encoder, effort, contact, force, imu) - for sensor in sensor_classes: - sensors_options = robot_options["sensors"][sensor.type] + for cls in (EncoderSensor, + EffortSensor, + ContactSensor, + ForceSensor, + ImuSensor): + sensors_options = robot_options["sensors"][cls.type] for sensor_options in sensors_options.values(): for name in ("delay", "jitter"): sensor_options[name] = sample( low=0.0, high=(self.std_ratio['sensors'] * - SENSOR_DELAY_SCALE[sensor.type]), + SENSOR_DELAY_SCALE[cls]), rg=self.np_random) for name in ( ("bias", SENSOR_BIAS_SCALE), ("noiseStd", SENSOR_NOISE_SCALE)): sensor_options[name] = sample( scale=(self.std_ratio['sensors'] * - SENSOR_NOISE_SCALE[sensor.type]), - shape=(len(sensor.fieldnames),), + SENSOR_NOISE_SCALE[cls]), + shape=(len( + cls.fieldnames),), # type: ignore[arg-type] rg=self.np_random) # Randomize the flexibility parameters @@ -361,7 +356,7 @@ def _force_external_profile(self, wrench[1] = F_PROFILE_SCALE * self._f_xy_profile[1](t) wrench[:2] *= self.std_ratio['disturbance'] - def has_terminated(self) -> Tuple[bool, bool]: + def has_terminated(self, info: InfoType) -> Tuple[bool, bool]: """Determine whether the episode is over. It terminates (`terminated=True`) under the following conditions: @@ -374,10 +369,12 @@ def has_terminated(self) -> Tuple[bool, bool]: - observation out-of-bounds - maximum simulation duration exceeded + :param info: Dictionary of extra information for monitoring. + :returns: terminated and truncated flags. """ # Call base implementation - terminated, truncated = super().has_terminated() + terminated, truncated = super().has_terminated(info) # Check if the agent has successfully solved the task if self._robot_state_q[2] < self._height_neutral * 0.5: @@ -385,10 +382,7 @@ def has_terminated(self) -> Tuple[bool, bool]: return terminated, truncated - def compute_reward(self, - terminated: bool, - truncated: bool, - info: InfoType) -> float: + def compute_reward(self, terminated: bool, info: InfoType) -> float: """Compute reward at current episode state. It computes the reward associated with each individual contribution @@ -409,7 +403,7 @@ def compute_reward(self, reward_dict['survival'] = 1.0 if 'energy' in reward_mixture_keys: - v_mot = self.robot.sensor_measurements[encoder.type][1] + _, v_mot = self.robot.sensor_measurements[EncoderSensor.type] command = self.robot_state.command power_consumption = np.sum(np.maximum(command * v_mot, 0.0)) power_consumption_rel = \ diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py index 17bef6054..da3ca4f99 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py @@ -1,15 +1,26 @@ # pylint: disable=missing-module-docstring from .manager import QuantityManager -from .generic import (AverageSpatialVelocityFrame, - EulerAnglesFrame) -from .locomotion import CenterOfMass, ZeroMomentPoint +from .generic import (FrameEulerAngles, + FrameXYZQuat, + StackedQuantity, + AverageFrameSpatialVelocity, + MaskedQuantity, + BinaryOpQuantity) +from .locomotion import (AverageOdometryVelocity, + CenterOfMass, + ZeroMomentPoint) __all__ = [ 'QuantityManager', - 'AverageSpatialVelocityFrame', - 'EulerAnglesFrame', + 'FrameEulerAngles', + 'FrameXYZQuat', + 'StackedQuantity', + 'AverageFrameSpatialVelocity', + 'MaskedQuantity', + 'BinaryOpQuantity', + 'AverageOdometryVelocity', 'CenterOfMass', 'ZeroMomentPoint', ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py index b94e2d855..a78159cc7 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py @@ -2,32 +2,560 @@ its topology (multiple or single branch, fixed or floating base...) and the application (locomotion, grasping...). """ +from copy import deepcopy +from collections import deque from functools import partial from dataclasses import dataclass -from typing import List, Dict, Set, Optional +from typing import ( + List, Dict, Set, Any, Optional, Protocol, Sequence, Tuple, TypeVar, Union, + Generic, Callable, runtime_checkable) +from typing_extensions import TypeAlias import numpy as np from jiminy_py.core import ( # pylint: disable=no-name-in-module - array_copyto, multi_array_copyto) + multi_array_copyto) import pinocchio as pin -from ..bases import InterfaceJiminyEnv, AbstractQuantity -from ..utils import fill, transforms_to_vector, matrix_to_rpy +from ..bases import ( + InterfaceJiminyEnv, InterfaceQuantity, AbstractQuantity, QuantityCreator, + QuantityEvalMode) +from ..utils import ( + fill, matrix_to_rpy, matrix_to_quat, quat_to_matrix, + quat_interpolate_middle) + + +EllipsisType: TypeAlias = Any # TODO: `EllipsisType` introduced in Python 3.10 + +ValueT = TypeVar('ValueT') +OtherValueT = TypeVar('OtherValueT') +YetAnotherValueT = TypeVar('YetAnotherValueT') + + +@runtime_checkable +class FrameQuantity(Protocol): + """Protocol that must be satisfied by all quantities associated with one + particular frame. + + This protocol is used when aggregating individual frame-level quantities + in a larger batch for computation vectorization on all frames at once. + Intermediate quantities managing these batches will make sure that all + their parents derive from one of the supported protocols, which includes + this one. + """ + frame_name: str + + +@runtime_checkable +class MultiFrameQuantity(Protocol): + """Protocol that must be satisfied by all quantities associated with + a particular set of frames for which the same batched intermediary + quantities must be computed. + + This protocol is involved in automatic computation vectorization. See + `FrameQuantity` documentation for details. + """ + frame_names: Sequence[str] + + +@dataclass(unsafe_hash=True) +class _MultiFramesRotationMatrix(AbstractQuantity[np.ndarray]): + """3D rotation matrix of the orientation of all frames involved in + quantities relying on it and are active since last reset of computation + tracking if shared cache is available, its parent otherwise. + + This quantity only provides a performance benefit when managed by some + `QuantityManager`. It is not supposed to be instantiated manually but use + as requirement of some other quantity for computation vectorization on all + frames at once. + + The data associated with each frame is exposed to the user as a batched 3D + contiguous array. The two first dimensions are rotation matrix elements, + while the last one are individual frames with the same ordering as + 'self.frame_names'. + + .. note:: + This quantity does not allow for specifying frames directly. There is + no way to get the orientation of multiple frames at once for now. + """ + + identifier: int + """Uniquely identify its parent type. + + This implies that quantities specifying `_MultiFramesRotationMatrix` as a + requirement will shared a unique batch with all the other ones having the + same type but not the others. This is essential to provide data access as a + batched ND contiguous array. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: InterfaceQuantity, + mode: QuantityEvalMode) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement. + :param mode: Desired mode of evaluation for this quantity. + """ + # Make sure that a parent has been specified + assert isinstance(parent, (FrameQuantity, MultiFrameQuantity)) + + # Set unique identifier based on parent type + self.identifier = hash(type(parent)) + + # Call base implementation + super().__init__( + env, parent, requirements={}, mode=mode, auto_refresh=False) + + # Initialize the ordered list of frame names + self.frame_names: Set[str] = set() + + # Store all rotation matrices at once + self._rot_mat_batch: np.ndarray = np.array([]) + + # Define proxy for slices of the batch storing all rotation matrices + self._rot_mat_slices: List[np.ndarray] = [] + + # Define proxy for the rotation matrices of all frames + self._rot_mat_list: List[np.ndarray] = [] + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Update the frame names based on the cache owners of this quantity. + # Note that only active quantities are considered for efficiency, which + # may change dynamically. Delegating this responsibility to cache + # owners may be possible but difficult to implement because + # `frame_names` must be cleared first before re-registering themselves, + # just in case of optimal computation graph has changed, not only once + # to avoid getting rid of quantities that just registered themselves. + # Nevertheless, whenever re-initializing this quantity to take into + # account changes of the active set must be decided by cache owners. + assert isinstance(self.parent, (FrameQuantity, MultiFrameQuantity)) + if isinstance(self.parent, FrameQuantity): + self.frame_names = {self.parent.frame_name} + else: + self.frame_names = set(self.parent.frame_names) + if self.has_cache: + for owner in self.cache.owners: + # We only consider active `_MultiFramesEulerAngles` instances + # instead of their parents. This is necessary because a derived + # quantity may feature `_MultiFramesEulerAngles` as requirement + # without actually relying on it depending on whether it is + # part of the optimal computation path at that time. + if owner.is_active(any_cache_owner=False): + assert isinstance( + owner.parent, (FrameQuantity, MultiFrameQuantity)) + if isinstance(owner.parent, FrameQuantity): + self.frame_names.add(owner.parent.frame_name) + else: + self.frame_names.union(owner.parent.frame_names) + + # Re-allocate memory as the number of frames is not known in advance. + # Note that Fortran memory layout (column-major) is used for speed up + # because it preserves contiguity when copying frame data. + nframes = len(self.frame_names) + self._rot_mat_batch = np.zeros((3, 3, nframes), order='F') + + # Refresh proxies + self._rot_mat_slices.clear() + self._rot_mat_list.clear() + for i, frame_name in enumerate(self.frame_names): + frame_index = self.pinocchio_model.getFrameId(frame_name) + rot_matrix = self.pinocchio_data.oMf[frame_index].rotation + self._rot_mat_slices.append(self._rot_mat_batch[..., i]) + self._rot_mat_list.append(rot_matrix) + + def refresh(self) -> np.ndarray: + # Copy all rotation matrices in contiguous buffer + multi_array_copyto(self._rot_mat_slices, self._rot_mat_list) + + # Return proxy directly without copy + return self._rot_mat_batch + + +@dataclass(unsafe_hash=True) +class _MultiFramesEulerAngles(AbstractQuantity[Dict[str, np.ndarray]]): + """Euler angles (Roll-Pitch-Yaw) of the orientation of all frames involved + in quantities relying on it and are active since last reset of computation + tracking if shared cache is available, its parent otherwise. + + It is not supposed to be instantiated manually but use internally by + `FrameEulerAngles`. See `_MultiFramesRotationMatrix` documentation. + + The orientation of all frames is exposed to the user as a dictionary whose + keys are the individual frame names. Internally, data are stored in batched + 2D contiguous array for efficiency. The first dimension gathers the 3 Euler + angles (roll, pitch, yaw), while the second one are individual frames with + the same ordering as 'self.frame_names'. + + The expected maximum speedup wrt computing Euler angles individually is + about x15, which is achieved asymptotically for more than 100 frames. It is + already x5 faster for 5 frames, x7 for 10 frames, and x9 for 20 frames. + """ + def __init__(self, + env: InterfaceJiminyEnv, + parent: "FrameEulerAngles", + mode: QuantityEvalMode) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: `FrameEulerAngles` instance from which this quantity is + a requirement. + :param mode: Desired mode of evaluation for this quantity. + """ + # Make sure that a suitable parent has been provided + assert isinstance(parent, FrameEulerAngles) + + # Initialize the ordered list of frame names. + # Note that this must be done BEFORE calling base `__init__`, otherwise + # `isinstance(..., (FrameQuantity, MultiFrameQuantity))` will fail. + self.frame_names: Set[str] = set() + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + rot_mat_batch=(_MultiFramesRotationMatrix, dict( + mode=mode))), + mode=mode, + auto_refresh=False) + + # Store Roll-Pitch-Yaw of all frames at once + self._rpy_batch: np.ndarray = np.array([]) + + # Mapping from frame name to individual Roll-Pitch-Yaw slices + self._rpy_map: Dict[str, np.ndarray] = {} + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Update the frame names based on the cache owners of this quantity + assert isinstance(self.parent, FrameEulerAngles) + self.frame_names = {self.parent.frame_name} + if self.has_cache: + for owner in self.cache.owners: + if owner.is_active(any_cache_owner=False): + assert isinstance(owner.parent, FrameEulerAngles) + self.frame_names.add(owner.parent.frame_name) + + # Re-allocate memory as the number of frames is not known in advance + nframes = len(self.frame_names) + self._rpy_batch = np.zeros((3, nframes), order='F') + + # Re-assign mapping from frame name to their corresponding data + self._rpy_map = dict(zip(self.frame_names, self._rpy_batch.T)) + + def refresh(self) -> Dict[str, np.ndarray]: + # Convert all rotation matrices at once to Roll-Pitch-Yaw + matrix_to_rpy(self.rot_mat_batch, self._rpy_batch) + + # Return proxy directly without copy + return self._rpy_map + + +@dataclass(unsafe_hash=True) +class FrameEulerAngles(InterfaceQuantity[np.ndarray]): + """Euler angles (Roll-Pitch-Yaw) of the orientation of a given frame in + world reference frame at the end of the agent step. + """ + + frame_name: str + """Name of the frame on which to operate. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + frame_name: str, + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param frame_name: Name of the frame on which to operate. + :param mode: Desired mode of evaluation for this quantity. + """ + # Backup some user argument(s) + self.frame_name = frame_name + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict(data=(_MultiFramesEulerAngles, dict(mode=mode))), + auto_refresh=False) + + def initialize(self) -> None: + # Check if the quantity is already active + was_active = self._is_active + + # Call base implementation. + # The quantity is now considered active at this point. + super().initialize() + + # Force re-initializing shared data if the active set has changed + if not was_active: + # Must reset the tracking for shared computation systematically, + # just in case the optimal computation path has changed to the + # point that relying on batched quantity is no longer relevant. + self.requirements["data"].reset(reset_tracking=True) + + def refresh(self) -> np.ndarray: + # Return a slice of batched data. + # Note that mapping from frame name to frame index in batched data + # cannot be pre-computed as it may changed dynamically. + return self.data[self.frame_name] + + +@dataclass(unsafe_hash=True) +class _MultiFramesXYZQuat(AbstractQuantity[Dict[str, np.ndarray]]): + """Vector representation (X, Y, Z, QuatX, QuatY, QuatZ, QuatW) of the + transform of all frames involved in quantities relying on it and are active + since last reset of computation tracking if shared cache is available, its + parent otherwise. + + It is not supposed to be instantiated manually but use internally by + `FrameXYZQuat`. See `_MultiFramesRotationMatrix` documentation. + + The transform of all frames is exposed to the user as a dictionary whose + keys are the individual frame names. Internally, data are stored in batched + 2D contiguous array for efficiency. The first dimension gathers the 6 + components (X, Y, Z, QuatX, QuatY, QuatZ, QuatW), while the second one are + individual frames with the same ordering as 'self.frame_names'. + """ + def __init__(self, + env: InterfaceJiminyEnv, + parent: "FrameXYZQuat", + mode: QuantityEvalMode) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: `FrameXYZQuat` instance from which this quantity + is a requirement. + :param mode: Desired mode of evaluation for this quantity. + """ + # Make sure that a suitable parent has been provided + assert isinstance(parent, FrameXYZQuat) + + # Initialize the ordered list of frame names + self.frame_names: Set[str] = set() + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + rot_mat_batch=(_MultiFramesRotationMatrix, dict( + mode=mode))), + mode=mode, + auto_refresh=False) + + # Define proxy for slices of the batch storing all translation vectors + self._translation_slices: List[np.ndarray] = [] + + # Define proxy for the translation vectors of all frames + self._translation_list: List[np.ndarray] = [] + + # Store XYZQuat of all frames at once + self._xyzquat_batch: np.ndarray = np.array([]) + + # Mapping from frame name to individual XYZQuat slices + self._xyzquat_map: Dict[str, np.ndarray] = {} + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Update the frame names based on the cache owners of this quantity + assert isinstance(self.parent, FrameXYZQuat) + self.frame_names = {self.parent.frame_name} + if self.has_cache: + for owner in self.cache.owners: + if owner.is_active(any_cache_owner=False): + assert isinstance(owner.parent, FrameXYZQuat) + self.frame_names.add(owner.parent.frame_name) + + # Re-allocate memory as the number of frames is not known in advance + nframes = len(self.frame_names) + self._xyzquat_batch = np.zeros((7, nframes), order='F') + + # Refresh proxies + self._translation_slices.clear() + self._translation_list.clear() + for i, frame_name in enumerate(self.frame_names): + frame_index = self.pinocchio_model.getFrameId(frame_name) + translation = self.pinocchio_data.oMf[frame_index].translation + self._translation_slices.append(self._xyzquat_batch[:3, i]) + self._translation_list.append(translation) + + # Re-assign mapping from frame name to their corresponding data + self._xyzquat_map = dict(zip(self.frame_names, self._xyzquat_batch.T)) + + def refresh(self) -> Dict[str, np.ndarray]: + # Copy all translations in contiguous buffer + multi_array_copyto(self._translation_slices, self._translation_list) + + # Convert all rotation matrices at once to XYZQuat representation + matrix_to_quat(self.rot_mat_batch, self._xyzquat_batch[-4:]) + + # Return proxy directly without copy + return self._xyzquat_map @dataclass(unsafe_hash=True) -class AverageSpatialVelocityFrame(AbstractQuantity[np.ndarray]): - """Average spatial velocity of a given frame at the end of an agent step. +class FrameXYZQuat(InterfaceQuantity[np.ndarray]): + """Vector representation (X, Y, Z, QuatX, QuatY, QuatZ, QuatW) of the + transform of a given frame in world reference frame at the end of the + agent step. + """ + + frame_name: str + """Name of the frame on which to operate. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + frame_name: str, + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param frame_name: Name of the frame on which to operate. + :param mode: Desired mode of evaluation for this quantity. + """ + # Backup some user argument(s) + self.frame_name = frame_name + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict(data=(_MultiFramesXYZQuat, dict(mode=mode))), + auto_refresh=False) + + def initialize(self) -> None: + # Check if the quantity is already active + was_active = self._is_active + + # Call base implementation + super().initialize() + + # Force re-initializing shared data if the active set has changed + if not was_active: + self.requirements["data"].reset(reset_tracking=True) + + def refresh(self) -> np.ndarray: + # Return a slice of batched data + return self.data[self.frame_name] + + +@dataclass(unsafe_hash=True) +class StackedQuantity(InterfaceQuantity[Tuple[ValueT, ...]]): + """Keep track of a given quantity over time by automatically stacking its + value once per environment step since last reset. + + .. note:: + A new entry is added to the stack right before evaluating the reward + and termination conditions. Internal simulation steps, observer and + controller updates are ignored. + """ + + quantity: InterfaceQuantity + """Base quantity whose value must be stacked over time since last reset. + """ + + num_stack: Optional[int] + """Maximum number of values that keep in memory before starting to discard + the oldest one (FIFO). None if unlimited. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + quantity: QuantityCreator[ValueT], + *, + num_stack: Optional[int] = None) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param quantity: Tuple gathering the class of the quantity whose values + must be stacked, plus all its constructor keyword- + arguments except environment 'env' and parent 'parent. + :param num_stack: Maximum number of values that keep in memory before + starting to discard the oldest one (FIFO). None if + unlimited. + """ + # Backup user arguments + self.num_stack = num_stack + + # Call base implementation + super().__init__(env, + parent, + requirements=dict(data=quantity), + auto_refresh=True) + + # Keep track of the quantity that must be stacked once instantiated + self.quantity = self.requirements["data"] + + # Allocate deque buffer + self._deque: deque = deque(maxlen=self.num_stack) + + # Keep track of the last time the quantity has been stacked + self._num_steps_prev = -1 + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Clear buffer + self._deque.clear() + + # Reset step counter + self._num_steps_prev = -1 + + def refresh(self) -> Tuple[ValueT, ...]: + # Append value to the queue only once per step and only if a simulation + # is running. Note that data must be deep-copied to make sure it does + # not get altered afterward. + if self.env.is_simulation_running: + num_steps = self.env.num_steps + if num_steps != self._num_steps_prev: + assert num_steps == self._num_steps_prev + 1 + self._deque.append(deepcopy(self.data)) + self._num_steps_prev += 1 + + # Return the whole stack as a tuple to preserve the integrity of the + # underlying container and make the API robust to internal changes. + return tuple(self._deque) + + +@dataclass(unsafe_hash=True) +class AverageFrameSpatialVelocity(InterfaceQuantity[np.ndarray]): + """Average spatial velocity of a given frame at the end of the agent step. The average spatial velocity is obtained by finite difference. More precisely, it is defined here as the ratio of the geodesic distance in SE3 - Lie group between the pose of the frame at the end of previous and current + Lie Group between the pose of the frame at the end of previous and current step over the time difference between them. Notably, under this definition, the linear average velocity jointly depends on rate of change of the translation and rotation of the frame, which may be undesirable in some cases. Alternatively, the double geodesic distance could be used instead to completely decouple the translation from the rotation. + + .. note:: + The local frame for which the velocity is expressed is defined as the + midpoint interpolation between the previous and current frame pose. + This definition is arbitrary, in a sense that any other point for an + interpolation ratio going from 0.0 (previous pose) to 1.0 (current + pose) would be equally valid. """ frame_name: str @@ -41,10 +569,11 @@ class AverageSpatialVelocityFrame(AbstractQuantity[np.ndarray]): def __init__(self, env: InterfaceJiminyEnv, - parent: Optional[AbstractQuantity], + parent: Optional[InterfaceQuantity], frame_name: str, - reference_frame: pin.ReferenceFrame = pin.LOCAL - ) -> None: + *, + reference_frame: pin.ReferenceFrame = pin.LOCAL, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: """ :param env: Base or wrapped jiminy environment. :param parent: Higher-level quantity from which this quantity is a @@ -54,6 +583,7 @@ def __init__(self, Whether the spatial velocity must be computed in local reference frame (aka 'pin.LOCAL') or re-aligned with world axes (aka 'pin.LOCAL_WORLD_ALIGNED'). + :param mode: Desired mode of evaluation for this quantity. """ # Make sure at requested reference frame is supported if reference_frame not in (pin.LOCAL, pin.LOCAL_WORLD_ALIGNED): @@ -65,7 +595,14 @@ def __init__(self, self.reference_frame = reference_frame # Call base implementation - super().__init__(env, parent, requirements={}) + super().__init__( + env, + parent, + requirements=dict(xyzquat_stack=(StackedQuantity, dict( + quantity=(FrameXYZQuat, dict( + frame_name=frame_name, mode=mode)), + num_stack=2))), + auto_refresh=False) # Define specialize difference operator on SE3 Lie group self._se3_diff = partial( @@ -75,12 +612,9 @@ def __init__(self, # Inverse step size self._inv_step_dt = 0.0 - # Pre-allocate memory to store current and previous frame pose vector - self._xyzquat_prev, self._xyzquat = ( - np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) for _ in range(2)) - - # Define proxy to the current frame pose (translation, rotation matrix) - self._pose = (np.zeros(3), np.eye(3)) + # Allocate memory for the average quaternion and rotation matrix + self._quat_mean = np.zeros(4) + self._rot_mat_mean = np.eye(3) # Pre-allocate memory for the spatial velocity self._v_spatial: np.ndarray = np.zeros(6) @@ -95,188 +629,204 @@ def initialize(self) -> None: # Compute inverse step size self._inv_step_dt = 1.0 / self.env.step_dt - # Refresh proxy to current frame pose - frame_index = self.pinocchio_model.getFrameId(self.frame_name) - transform = self.pinocchio_data.oMf[frame_index] - self._pose = (transform.translation, transform.rotation) - # Re-initialize pre-allocated buffers - transforms_to_vector((self._pose,), self._xyzquat) - array_copyto(self._xyzquat_prev, self._xyzquat) fill(self._v_spatial, 0) def refresh(self) -> np.ndarray: - # Convert current transform to (XYZ, Quat) convention - transforms_to_vector((self._pose,), self._xyzquat) + # Fetch previous and current XYZQuat representation of frame transform. + # It will raise an exception if not enough data is available at this + # point. This should never occur in practice as it will be fine at + # the end of the first step already, before the reward and termination + # conditions are evaluated. + xyzquat_prev, xyzquat = self.xyzquat_stack # Compute average frame velocity in local frame since previous step - self._v_spatial[:] = self._se3_diff(self._xyzquat_prev, self._xyzquat) + self._v_spatial[:] = self._se3_diff(xyzquat_prev, xyzquat) self._v_spatial *= self._inv_step_dt # Translate local velocity to world frame if self.reference_frame == pin.LOCAL_WORLD_ALIGNED: - # TODO: x2 speedup can be expected using `np.dot` with `nb.jit` - _, rot_mat = self._pose - self._v_lin_ang[:] = rot_mat @ self._v_lin_ang - - # Backup current frame pose - array_copyto(self._xyzquat_prev, self._xyzquat) + # Define world frame as the "middle" between prev and next pose. + # The orientation difference has an effect on the translation + # difference, but not the other way around. Here, we only care + # about the middle rotation, so we can consider SO3 Lie Group + # algebra instead of SE3. + quat_interpolate_middle( + xyzquat_prev[-4:], xyzquat[-4:], self._quat_mean) + quat_to_matrix(self._quat_mean, self._rot_mat_mean) + + # TODO: x2 speedup can be expected using `np.dot` with `nb.jit` + self._v_lin_ang[:] = self._rot_mat_mean @ self._v_lin_ang return self._v_spatial @dataclass(unsafe_hash=True) -class _BatchEulerAnglesFrame(AbstractQuantity[Dict[str, np.ndarray]]): - """Euler angles (Roll-Pitch-Yaw) representation of the orientation of all - frames involved in quantities relying on it and are active since last reset - of computation tracking if shared cache is available, its parent otherwise. - - This quantity only provides a performance benefit when managed by some - `QuantityManager`. It is not supposed to be instantiated manually but use - internally by `EulerAnglesFrame` as a requirement for vectorization of - computations for all frames at once. +class MaskedQuantity(InterfaceQuantity[np.ndarray]): + """Extract elements from a given quantity whose value is a N-dimensional + array along an axis. + + Elements will be extract by copy unless the indices of the elements to + extract to be written equivalently by a slice (ie they are evenly spaced), + and the array can be flattened while preserving memory contiguity if 'axis' + is `None`. + """ - The orientation of all frames is exposed to the user as a dictionary whose - keys are the individual frame names. Internally, data are stored in batched - 2D contiguous array for efficiency. The first dimension are the 3 Euler - angles (roll, pitch, yaw) components, while the second one are individual - frames with the same ordering as 'self.frame_names'. + quantity: InterfaceQuantity + """Base quantity whose elements must be extracted. + """ - The expected maximum speedup wrt computing Euler angles individually is - about x15, which is achieved asymptotically for more than 100 frames. It is - already x5 faster for 5 frames, x7 for 10 frames, and x9 for 20 frames. + indices: Tuple[int, ...] + """Indices of the elements to extract. + """ - .. note:: - This quantity does not allow for specifying frames directly. There is - no way to get the orientation of multiple frames at once for now. + axis: Optional[int] + """Axis over which to extract elements. `None` to consider flattened array. """ + def __init__(self, env: InterfaceJiminyEnv, - parent: Optional[AbstractQuantity]) -> None: + parent: Optional[InterfaceQuantity], + quantity: QuantityCreator[np.ndarray], + key: Union[Sequence[int], Sequence[bool]], + *, + axis: Optional[int] = None) -> None: """ :param env: Base or wrapped jiminy environment. :param parent: Higher-level quantity from which this quantity is a requirement if any, `None` otherwise. - :param frame_name: Name of the frame on which to operate. + :param quantity: Tuple gathering the class of the quantity whose values + must be extracted, plus all its constructor keyword- + arguments except environment 'env' and parent 'parent. + :param key: Sequence of indices or boolean mask that will be used to + extract elements from the quantity along one axis. + :param axis: Axis over which to extract elements. `None` to consider + flattened array. + Optional: `None` by default. """ - # Call base implementation - super().__init__(env, parent, requirements={}) - - # Initialize the ordered list of frame names - self.frame_names: Set[str] = set() - - # Store all rotation matrices at once - self._rot_mat_batch: np.ndarray = np.array([]) - - # Define proxy for slices of the batch storing all rotation matrices - self._rot_mat_views: List[np.ndarray] = [] - - # Define proxy for the rotation matrices of all frames - self._rot_mat_list: List[np.ndarray] = [] + # Check if indices or boolean mask has been provided + if all(isinstance(e, bool) for e in key): + key = tuple(np.flatnonzero(key)) + elif not all(isinstance(e, int) for e in key): + raise ValueError( + "Argument 'key' invalid. It must either be a boolean mask, or " + "a sequence of indices.") + + # Backup user arguments + self.indices = tuple(key) + self.axis = axis + + # Make sure that at least one index must be extracted + if not self.indices: + raise ValueError( + "No indices to extract from quantity. Data would be empty.") + + # Check if the indices are evenly spaced + self._slices: Tuple[Union[slice, EllipsisType], ...] = () + stride: Optional[int] = None + if len(self.indices) == 1: + stride = 1 + if len(self.indices) > 1 and all(e >= 0 for e in self.indices): + spacing = np.unique(np.diff(self.indices)) + if spacing.size == 1: + stride = spacing[0] + if stride is not None: + slice_ = slice(self.indices[0], self.indices[-1] + 1, stride) + if axis is None: + self._slices = (slice_,) + elif axis > 0: + self._slices = (*((slice(None),) * axis), slice_) + else: + self._slices = ( + Ellipsis, slice_, *((slice(None),) * (- axis - 1))) - # Store Roll-Pitch-Yaw of all frames at once - self._rpy_batch: np.ndarray = np.array([]) + # Call base implementation + super().__init__(env, + parent, + requirements=dict(data=quantity), + auto_refresh=False) - # Mapping from frame name to individual Roll-Pitch-Yaw slices - self._rpy_map: Dict[str, np.ndarray] = {} + # Keep track of the quantity from which data must be extracted + self.quantity = self.requirements["data"] def initialize(self) -> None: # Call base implementation super().initialize() - # Update the list of frame names based on its cache owner list. - # Note that only active quantities are shared for efficiency. The list - # of active quantity may change dynamically. Re-initializing the class - # to take into account changes of the active set must be decided by - # `EulerAnglesFrame`. - assert isinstance(self.parent, EulerAnglesFrame) - self.frame_names = {self.parent.frame_name} - if self.cache: - for owner in self.cache.owners: - # We only consider active instances of `_BatchEulerAnglesFrame` - # instead of their corresponding parent `EulerAnglesFrame`. - # This is necessary because a derived quantity may feature - # `_BatchEulerAnglesFrame` as a requirement without actually - # relying on it depending on whether it is part of the optimal - # computation path at the time being or not. - if owner.is_active(any_cache_owner=False): - assert isinstance(owner.parent, EulerAnglesFrame) - self.frame_names.add(owner.parent.frame_name) - - # Re-allocate memory as the number of frames is not known in advance. - # Note that Fortran memory layout (column-major) is used for speed up - # because it preserves contiguity when copying frame data. - nframes = len(self.frame_names) - self._rot_mat_batch = np.zeros((3, 3, nframes), order='F') - self._rpy_batch = np.zeros((3, nframes)) - - # Refresh proxies - self._rot_mat_views.clear() - self._rot_mat_list.clear() - for i, frame_name in enumerate(self.frame_names): - frame_index = self.pinocchio_model.getFrameId(frame_name) - rot_matrix = self.pinocchio_data.oMf[frame_index].rotation - self._rot_mat_views.append(self._rot_mat_batch[..., i]) - self._rot_mat_list.append(rot_matrix) - - # Re-assign mapping from frame name to their corresponding data - self._rpy_map = dict(zip(self.frame_names, self._rpy_batch.T)) + def refresh(self) -> np.ndarray: + # Extract elements from quantity + if not self._slices: + # Note that `take` is faster than classical advanced indexing via + # `operator[]` (`__getitem__`) because the latter is more generic. + # Notably, `operator[]` supports boolean mask but `take` does not. + return self.data.take(self.indices, axis=self.axis) + if self.axis is None: + # `reshape` must be used instead of `flat` to get a view that can + # be sliced without copy. + return self.data.reshape((-1,))[self._slices] + return self.data[self._slices] - def refresh(self) -> Dict[str, np.ndarray]: - # Copy all rotation matrices in contiguous buffer - multi_array_copyto(self._rot_mat_views, self._rot_mat_list) - # Convert all rotation matrices at once to Roll-Pitch-Yaw - matrix_to_rpy(self._rot_mat_batch, self._rpy_batch) +@dataclass(unsafe_hash=True) +class BinaryOpQuantity(InterfaceQuantity[ValueT], + Generic[ValueT, OtherValueT, YetAnotherValueT]): + """Apply a given binary operator between two quantities. - # Return proxy directly without copy - return self._rpy_map + This quantity is mainly useful for computing the error between the value of + a given quantity evaluated at the current simulation state and the state of + at the current simulation time for the reference trajectory being selected. + """ + quantity_left: InterfaceQuantity[OtherValueT] + """Left-hand side quantity that will be forwarded to the binary operator. + """ -@dataclass(unsafe_hash=True) -class EulerAnglesFrame(AbstractQuantity[np.ndarray]): - """Euler angles (Roll-Pitch-Yaw) representation of the orientation of a - given frame in world reference frame at the end of an agent step. + quantity_right: InterfaceQuantity[YetAnotherValueT] + """Right-hand side quantity that will be forwarded to the binary operator. """ - frame_name: str - """Name of the frame on which to operate. + op: Callable[[OtherValueT, YetAnotherValueT], ValueT] + """Callable taking right- and left-hand side quantities as input argument. """ def __init__(self, env: InterfaceJiminyEnv, - parent: Optional[AbstractQuantity], - frame_name: str, + parent: Optional[InterfaceQuantity], + quantity_left: QuantityCreator[OtherValueT], + quantity_right: QuantityCreator[YetAnotherValueT], + op: Callable[[OtherValueT, YetAnotherValueT], ValueT] ) -> None: """ :param env: Base or wrapped jiminy environment. :param parent: Higher-level quantity from which this quantity is a requirement if any, `None` otherwise. - :param frame_name: Name of the frame on which to operate. + :param quantity_left: Tuple gathering the class of the quantity that + must be passed to left-hand side of the binary + operator, plus all its constructor keyword- + arguments except environment 'env' and parent + 'parent. + :param quantity_right: Quantity that must be passed to right-hand side + of the binary operator as a tuple + (class, keyword-arguments). See `quantity_left` + argument for details. + :param op: Any callable taking the right- and left-hand side quantities + as input argument. For example `operator.sub` to compute the + difference. """ # Backup some user argument(s) - self.frame_name = frame_name + self.op = op # Call base implementation super().__init__( - env, parent, requirements={"data": (_BatchEulerAnglesFrame, {})}) - - def initialize(self) -> None: - # Check if the quantity is already active - was_active = self._is_active - - # Call base implementation. - # The quantity is now considered active at this point. - super().initialize() - - # Force re-initializing shared data if the active set has changed - if not was_active: - # Must reset the tracking for shared computation systematically, - # just in case the optimal computation path has changed. - self.requirements["data"].reset(reset_tracking=True) - - def refresh(self) -> np.ndarray: - # Return a slice of batched data. - # Note that mapping from frame name to frame index in batched data - # cannot be pre-computed as it may changed dynamically. - return self.data[self.frame_name] + env, + parent, + requirements=dict( + value_left=quantity_left, value_right=quantity_right), + auto_refresh=False) + + # Keep track of the left- and right-hand side quantities for hashing + self.quantity_left = self.requirements["value_left"] + self.quantity_right = self.requirements["value_right"] + + def refresh(self) -> ValueT: + return self.op(self.value_left, self.value_right) diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py index 5cf397166..af8ddd16c 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py @@ -6,9 +6,52 @@ import numpy as np import pinocchio as pin -from ..bases import InterfaceJiminyEnv, AbstractQuantity +from ..bases import ( + InterfaceJiminyEnv, InterfaceQuantity, AbstractQuantity, QuantityEvalMode) from ..utils import fill +from ..quantities import MaskedQuantity, AverageFrameSpatialVelocity + + +@dataclass(unsafe_hash=True) +class AverageOdometryVelocity(InterfaceQuantity[np.ndarray]): + """Average odometry velocity in local-world-aligned frame at the end of the + agent step. + + The odometry pose fully specifies the position of the robot in 2D world + plane. As such, it comprises the linear translation (X, Y) and the angular + velocity around Z axis (namely rate of change of Yaw Euler angle). The + average spatial velocity is obtained by finite difference. See + `AverageFrameSpatialVelocity` documentation for details. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param mode: Desired mode of evaluation for this quantity. + """ + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + data=(MaskedQuantity, dict( + quantity=(AverageFrameSpatialVelocity, dict( + frame_name="root_joint", + reference_frame=pin.LOCAL_WORLD_ALIGNED, + mode=mode)), + key=(0, 1, 5)))), + auto_refresh=False) + + def refresh(self) -> np.ndarray: + return self.data + @dataclass(unsafe_hash=True) class CenterOfMass(AbstractQuantity[np.ndarray]): @@ -23,21 +66,24 @@ class CenterOfMass(AbstractQuantity[np.ndarray]): def __init__( self, env: InterfaceJiminyEnv, - parent: Optional[AbstractQuantity], - kinematic_level: pin.KinematicLevel = pin.POSITION - ) -> None: + parent: Optional[InterfaceQuantity], + *, + kinematic_level: pin.KinematicLevel = pin.POSITION, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: """ :param env: Base or wrapped jiminy environment. :param parent: Higher-level quantity from which this quantity is a requirement if any, `None` otherwise. :para kinematic_level: Desired kinematic level, ie position, velocity or acceleration. + :param mode: Desired mode of evaluation for this quantity. """ # Backup some user argument(s) self.kinematic_level = kinematic_level # Call base implementation - super().__init__(env, parent, requirements={}) + super().__init__( + env, parent, requirements={}, mode=mode, auto_refresh=False) # Pre-allocate memory for the CoM quantity self._com_data: np.ndarray = np.array([]) @@ -46,6 +92,13 @@ def initialize(self) -> None: # Call base implementation super().initialize() + # Make sure that the state is consistent with required kinematic level + state = self.state + if ((self.kinematic_level == pin.ACCELERATION and state.a is None) or + (self.kinematic_level >= pin.VELOCITY and state.v is None)): + raise RuntimeError( + "State data inconsistent with required kinematic level") + # Refresh CoM quantity proxy based on kinematic level if self.kinematic_level == pin.POSITION: self._com_data = self.pinocchio_data.com[0] @@ -56,10 +109,11 @@ def initialize(self) -> None: def refresh(self) -> np.ndarray: # Jiminy does not compute the CoM acceleration automatically - if self.kinematic_level == pin.ACCELERATION: + if (self.mode == QuantityEvalMode.TRUE and + self.kinematic_level == pin.ACCELERATION): pin.centerOfMass(self.pinocchio_model, self.pinocchio_data, - self.kinematic_level) + pin.ACCELERATION) # Return proxy directly without copy return self._com_data @@ -78,18 +132,28 @@ class ZeroMomentPoint(AbstractQuantity[np.ndarray]): """ def __init__(self, env: InterfaceJiminyEnv, - parent: Optional[AbstractQuantity]) -> None: + parent: Optional[InterfaceQuantity], + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: """ :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param mode: Desired mode of evaluation for this quantity. """ # Call base implementation - super().__init__(env, parent, requirements={"com": (CenterOfMass, {})}) + super().__init__( + env, + parent, + requirements=dict(com=(CenterOfMass, dict(mode=mode))), + mode=mode, + auto_refresh=False) # Weight of the robot self._robot_weight: float = -1 # Proxy for the derivative of the spatial centroidal momentum - self.dhg: Tuple[np.ndarray, np.ndarray] = (np.ndarray([]),) * 2 + self.dhg: Tuple[np.ndarray, np.ndarray] = (np.array([]),) * 2 # Pre-allocate memory for the ZMP self._zmp = np.zeros(2) @@ -98,6 +162,11 @@ def initialize(self) -> None: # Call base implementation super().initialize() + # Make sure that the state is consistent with required kinematic level + if (self.state.v is None or self.state.a is None): + raise RuntimeError( + "State data inconsistent with required kinematic level") + # Compute the weight of the robot gravity = abs(self.pinocchio_model.gravity.linear[2]) robot_mass = self.pinocchio_data.mass[0] @@ -121,5 +190,4 @@ def refresh(self) -> np.ndarray: if abs(f_z) > np.finfo(np.float32).eps: self._zmp[0] -= (dhg_angular[1] + dhg_linear[0] * com[2]) / f_z self._zmp[1] += (dhg_angular[0] - dhg_linear[1] * com[2]) / f_z - return self._zmp diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/manager.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/manager.py index 44b14a10d..56659cac5 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/manager.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/manager.py @@ -9,21 +9,28 @@ quantities in a large batch to leverage vectorization of math instructions. """ from collections.abc import MutableMapping -from typing import Any, Dict, List, Tuple, Iterator, Type +from typing import Any, Dict, List, Tuple, Iterator, Literal, Type -from ..bases import InterfaceJiminyEnv, AbstractQuantity, SharedCache +from jiminy_py.dynamics import Trajectory - -QuantityCreator = Tuple[Type[AbstractQuantity], Dict[str, Any]] +from ..bases import ( + QuantityCreator, InterfaceJiminyEnv, InterfaceQuantity, SharedCache, + DatasetTrajectoryQuantity) class QuantityManager(MutableMapping): """This class centralizes the evaluation of all quantities involved in - reward or termination conditions evaluation to redundant and unnecessary - computations. + reward components or termination conditions evaluation to redundant and + unnecessary computations. It is responsible for making sure all quantities are evaluated on the same - environment, and internal buffers are re-initialized whenever necessary. + environment, and internal buffers are re-initialized whenever necessary. It + also manages a dataset of trajectories synchronized between all managed + quantities. These trajectories are involves in computation of quantities + deriving from `AbstractQuantity` for which the mode of evaluation is set to + `QuantityEvalMode.REFERENCE`. This dataset is initially empty. Trying to + evaluate quantities involving a reference trajectory without adding and + selecting one beforehand would raise an exception. .. note:: Individual quantities can be accessed either as instance properties or @@ -38,16 +45,22 @@ def __init__(self, env: InterfaceJiminyEnv) -> None: # Backup user argument(s) self.env = env - # List of already instantiated quantities to manager - self._quantities: Dict[str, AbstractQuantity] = {} + # List of instantiated quantities to manager + self.registry: Dict[str, InterfaceQuantity] = {} # Initialize shared caches for all managed quantities. # Note that keys are not quantities directly but pairs (class, hash). # This is necessary because using a quantity as key directly would - # prevent its garbage collection and break automatic reset of + # prevent its garbage collection, hence breaking automatic reset of # computation tracking for all quantities sharing its cache. self._caches: Dict[ - Tuple[Type[AbstractQuantity], int], SharedCache] = {} + Tuple[Type[InterfaceQuantity], int], SharedCache] = {} + + # Instantiate trajectory database. + # Note that this quantity is not added to the global registry to avoid + # exposing directly to the user. This way, it cannot be deleted. + self._trajectory_dataset = self._build_quantity( + (DatasetTrajectoryQuantity, {})) def reset(self, reset_tracking: bool = False) -> None: """Consider that all managed quantity must be re-initialized before @@ -59,7 +72,7 @@ def reset(self, reset_tracking: bool = False) -> None: :param reset_tracking: Do not consider any quantity as active anymore. Optional: False by default. """ - for quantity in self._quantities.values(): + for quantity in self.registry.values(): quantity.reset(reset_tracking) def clear(self) -> None: @@ -74,6 +87,44 @@ def clear(self) -> None: for cache in self._caches.values(): cache.reset() + def add_trajectory(self, name: str, trajectory: Trajectory) -> None: + """Add a new reference trajectory to the database synchronized between + all managed quantities. + + :param name: Desired name of the trajectory. It must be unique. If a + trajectory with the exact same name already exists, then + it must be discarded first, so as to prevent silently + overwriting it by mistake. + :param trajectory: Trajectory instance to register. + """ + self._trajectory_dataset.add(name, trajectory) + + def discard_trajectory(self, name: str) -> None: + """Discard a trajectory from the database synchronized between all + managed quantities. + + :param name: Name of the trajectory to discard. + """ + self._trajectory_dataset.discard(name) + + def select_trajectory(self, + name: str, + mode: Literal['raise', 'wrap', 'clip'] = 'raise' + ) -> None: + """Select an existing trajectory from the database shared synchronized + all managed quantities. + + .. note:: + There is no way to select a different reference trajectory for + individual quantities at the time being. + + :param name: Name of the trajectory to select. + :param mode: Specifies how to deal with query time of are out of the + time interval of the trajectory. See `Trajectory.get` + documentation for details. + """ + self._trajectory_dataset.select(name, mode) + def __getattr__(self, name: str) -> Any: """Get access managed quantities as first-class properties, rather than dictionary-like values through `__getitem__`. @@ -95,27 +146,28 @@ def __dir__(self) -> List[str]: It is mainly used by autocomplete feature of Ipython. It is overloaded to get consistent autocompletion wrt `getattr`. """ - return [*super().__dir__(), *self._quantities.keys()] + return [*super().__dir__(), *self.registry.keys()] - def __setitem__(self, - name: str, - quantity_creator: QuantityCreator) -> None: - """Instantiate a new top-level quantity to manage for now on. + def _build_quantity( + self, quantity_creator: QuantityCreator) -> InterfaceQuantity: + """Instantiate a quantity sharing caching with all identical quantities + that has been instantiated previously. + + .. note:: + This method is not meant to be called manually. It is used + internally for instantiating new quantities sharing cache with all + identical instances that has been instantiated previously by this + manager in particular. This method is not responsible for keeping + track of the new quantity and exposing it to the user by adding it + to the global registry of the manager afterward. :param name: Desired name of the quantity after instantiation. It will raise an exception if another quantity with the exact same name exists. :param quantity_creator: Tuple gathering the class of the new quantity to manage plus its keyword-arguments except - the environment 'env' as a dictionary. + environment and parent as a dictionary. """ - # Make sure that no quantity with the same name is already managed to - # avoid silently overriding quantities being managed in user's back. - if name in self._quantities: - raise KeyError( - "A quantity with the exact same name already exists. Please " - "delete it first before adding a new one.") - # Instantiate the new quantity quantity_cls, quantity_kwargs = quantity_creator top_quantity = quantity_cls(self.env, None, **(quantity_kwargs or {})) @@ -128,15 +180,39 @@ def __setitem__(self, quantity.cache = self._caches.setdefault(key, SharedCache()) quantities_all += quantity.requirements.values() - # Add it to the map of already managed quantities - self._quantities[name] = top_quantity + return top_quantity + + def __setitem__(self, + name: str, + quantity_creator: QuantityCreator) -> None: + """Instantiate new top-level quantity that will be managed for now on. + + :param name: Desired name of the quantity after instantiation. It will + raise an exception if another quantity with the exact same + name exists. + :param quantity_creator: Tuple gathering the class of the new quantity + to manage plus its keyword-arguments except + environment and parent as a dictionary. + """ + # Make sure that no quantity with the same name is already managed to + # avoid silently overriding quantities being managed in user's back. + if name in self.registry: + raise KeyError( + "A quantity with the exact same name already exists. Please " + "delete it first before adding a new one.") + + # Instantiate new quantity + quantity = self._build_quantity(quantity_creator) + + # Add it to the global registry of already managed quantities + self.registry[name] = quantity def __getitem__(self, name: str) -> Any: """Get the evaluated value of a given quantity. :param name: Name of the quantity for which to fetch the current value. """ - return self._quantities[name].get() + return self.registry[name].get() def __delitem__(self, name: str) -> None: """Stop managing a quantity that is no longer relevant. @@ -151,14 +227,14 @@ def __delitem__(self, name: str) -> None: :param name: Name of the managed quantity to be discarded. It will raise an exception if the specified name does not exists. """ - del self._quantities[name] + del self.registry[name] def __iter__(self) -> Iterator[str]: """Iterate over names of managed quantities. """ - return iter(self._quantities) + return iter(self.registry) def __len__(self) -> int: """Number of quantities being managed. """ - return len(self._quantities) + return len(self.registry) diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/utils/__init__.py index 32b108c84..19ebdaa32 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/__init__.py @@ -11,6 +11,7 @@ quat_to_yaw_cos_sin, quat_multiply, quat_average, + quat_interpolate_middle, rpy_to_matrix, rpy_to_quat, transforms_to_vector, @@ -40,7 +41,10 @@ get_fieldnames, register_variables, sample) -from .pipeline import build_pipeline, load_pipeline +from .pipeline import (save_trajectory_to_hdf5, + load_trajectory_from_hdf5, + build_pipeline, + load_pipeline) __all__ = [ @@ -55,6 +59,7 @@ 'quat_to_yaw_cos_sin', 'quat_multiply', 'quat_average', + 'quat_interpolate_middle', 'rpy_to_matrix', 'rpy_to_quat', 'transforms_to_vector', @@ -84,6 +89,8 @@ 'is_nan', 'get_fieldnames', 'register_variables', + 'save_trajectory_to_hdf5', + 'load_trajectory_from_hdf5', 'build_pipeline', 'load_pipeline' ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py index 9721cd314..ffa47d52b 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py @@ -522,3 +522,38 @@ def quat_average(quat: np.ndarray, q_flat = q_perm.reshape((*q_perm.shape[:-len(axis)], -1)) _, eigvec = np.linalg.eigh(q_flat @ np.swapaxes(q_flat, -1, -2)) return np.moveaxis(eigvec[..., -1], -1, 0) + + +@nb.jit(nopython=True, cache=True, fastmath=True) +def quat_interpolate_middle(quat1: np.ndarray, + quat2: np.ndarray, + out: Optional[np.ndarray] = None) -> np.ndarray: + """Compute the midpoint interpolation between two batches of quaternions + [qx, qy, qz, qw]. + + The midpoint interpolation of two quaternion is defined as the integration + of half the difference between them, starting from the first one, ie + `q_mid = integrate(q1, 0.5 * difference(q1, d2))`, which is a special case + of the `slerp` method (spherical linear interpolation) for `alpha=0.5`. + + For the midpoint in particular, one can show that the middle quaternion is + simply normalized sum of the previous and next quaternions. + + :param quat1: First batch of quaternions as a N-dimensional array whose + first dimension gathers the 4 quaternion coordinates. + :param quat2: Second batch of quaternions as a N-dimensional array. + :param out: A pre-allocated array into which the result is stored. If not + provided, a new array is freshly-allocated, which is slower. + """ + assert quat1.ndim >= 1 and quat1.shape == quat2.shape + if out is None: + out_ = np.empty((4, *quat1.shape[1:])) + else: + assert out.shape == (4, *quat1.shape[1:]) + out_ = out + + dot = np.sum(quat1 * quat2, axis=0) + dot_ = dot if quat1.ndim == 1 else np.expand_dims(dot, axis=0) + out_[:] = (quat1 + np.sign(dot_) * quat2) / np.sqrt(2 * (1 + np.abs(dot_))) + + return out_ diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py b/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py index 5012c024f..15c58a99b 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py @@ -10,90 +10,205 @@ import json import pathlib from pydoc import locate +from dataclasses import asdict from functools import partial from typing import ( - Dict, Any, Optional, Union, Type, Sequence, Callable, TypedDict) + Dict, Any, Optional, Union, Type, Sequence, Callable, TypedDict, Literal, + cast) +import h5py import toml +import numpy as np import gymnasium as gym +import jiminy_py.core as jiminy +from jiminy_py.dynamics import State, Trajectory + from ..bases import (InterfaceJiminyEnv, InterfaceBlock, BaseControllerBlock, BaseObserverBlock, BasePipelineWrapper, ObservedJiminyEnv, - ControlledJiminyEnv) + ControlledJiminyEnv, + ComposedJiminyEnv, + AbstractReward, + BaseQuantityReward, + BaseMixtureReward) from ..envs import BaseJiminyEnv +class RewardConfig(TypedDict, total=False): + """Store information required for instantiating a given reward. + + Specifically, it is a dictionary comprising the class of the reward, which + must derive from `AbstractReward`, and optionally some keyword-arguments + that must be passed to its corresponding constructor. + """ + + cls: Union[Type[AbstractReward], str] + """Reward class type. + + .. note:: + Both class type or fully qualified dotted path are supported. + """ + + kwargs: Dict[str, Any] + """Environment constructor keyword-arguments. + + This attribute can be omitted. + """ + + +class TrajectoriesConfig(TypedDict, total=False): + """Store information required for adding a database of reference + trajectories to the environment. + + Specifically, it is a dictionary comprising a set of named trajectories as + a dictionary whose keys are the name of the trajectories and values are + either the trajectory itself or the path of a file storing its dump in HDF5 + format, the name of the selected trajectory, and its interpolation mode. + """ + + dataset: Dict[str, Union[str, Trajectory]] + """Set of named trajectories as a dictionary. + + .. note:: + Both `Trajectory` objects or path (absolute or relative) are supported. + """ + + name: str + """Name of the selected trajectory if any. + + This attribute can be omitted. + """ + + mode: Literal['raise', 'wrap', 'clip'] + """Interpolation mode of the selected trajectory if any. + + This attribute can be omitted. + """ + + class EnvConfig(TypedDict, total=False): + """Store information required for instantiating a given base environment + and compose it with some additional reward components and termination + conditions. + + Specifically, it is a dictionary comprising the class of the base + environment, which must derive from `BaseJiminyEnv`, optionally some + keyword-arguments that must be passed to its corresponding constructor, and + eventually the configuration of some additional reward with which to + compose the base environment. + """ + + cls: Union[Type[BaseJiminyEnv], str] """Environment class type. .. note:: Both class type or fully qualified dotted path are supported. """ - cls: Union[Type[BaseJiminyEnv], str] + kwargs: Dict[str, Any] """Environment constructor default arguments. This attribute can be omitted. """ - kwargs: Dict[str, Any] + + reward: RewardConfig + """Reward configuration. + + This attribute can be omitted. + """ + + trajectories: TrajectoriesConfig + """Reference trajectories configuration. + + This attribute can be omitted. + """ class BlockConfig(TypedDict, total=False): + """Store information required for instantiating a given observation or + control block. + + Specifically, it is a dictionary comprising the class of the block, which + must derive from `BaseControllerBlock` or `BaseObserverBlock`, and + optionally some keyword-arguments that must be passed to its corresponding + constructor. + """ + + cls: Union[Type[BaseControllerBlock], Type[BaseObserverBlock], str] """Block class type. If must derive from `BaseControllerBlock` for controller blocks or from `BaseObserverBlock` for observer blocks. .. note:: Both class type or fully qualified dotted path are supported. """ - cls: Union[Type[BaseControllerBlock], Type[BaseObserverBlock], str] + kwargs: Dict[str, Any] """Block constructor default arguments. This attribute can be omitted. """ - kwargs: Dict[str, Any] class WrapperConfig(TypedDict, total=False): + """Store information required for instantiating a given environment + pipeline wrapper. + + Specifically, it is a dictionary comprising the class of the wrapper, which + must derive from `BasePipelineWrapper`, and optionally some + keyword-arguments that must be passed to its corresponding constructor. + """ + + cls: Union[Type[BasePipelineWrapper], str] """Wrapper class type. .. note:: Both class type or fully qualified dotted path are supported. """ - cls: Union[Type[BasePipelineWrapper], str] + kwargs: Dict[str, Any] """Wrapper constructor default arguments. This attribute can be omitted. """ - kwargs: Dict[str, Any] class LayerConfig(TypedDict, total=False): - """Block constructor default arguments. + """Store information required for instantiating a given environment + pipeline layer, ie either a wrapper, or the combination of an observer / + controller block with its corresponding wrapper. + + Specifically, it is a dictionary comprising the configuration of the block + if any, and optionally the configuration of the reward. It is generally + sufficient to specify either one or the other. See the documentation of the + both fields for details. + """ + + block: BlockConfig + """Block configuration. This attribute can be omitted. If so, then 'wrapper_cls' must be specified and must not require any block. Typically, it happens when the wrapper is not doing any computation on its own but just transforming the action or observation, e.g. stacking observation frames. """ - block: Optional[BlockConfig] + wrapper: WrapperConfig """Wrapper configuration. This attribute can be omitted. If so, then 'block' must be specified and must this block must be associated with a unique wrapper type to allow for automatic type inference. It works with any observer and controller block. """ - wrapper: WrapperConfig def build_pipeline(env_config: EnvConfig, - layers_config: Sequence[LayerConfig] + layers_config: Sequence[LayerConfig], + *, + root_path: Optional[Union[str, pathlib.Path]] = None ) -> Callable[..., InterfaceJiminyEnv]: """Wrap together an environment inheriting from `BaseJiminyEnv` with any number of layers, as a unified pipeline environment class inheriting from @@ -107,6 +222,108 @@ def build_pipeline(env_config: EnvConfig, lowest level layer to the highest, each element corresponding to the configuration of a individual layer, as a dict of type `LayerConfig`. """ + # Define helper to sanitize reward configuration + def sanitize_reward_config(reward_config: RewardConfig) -> None: + """Sanitize reward configuration in-place. + + :param reward_config: Configuration of the reward, as a dict of type + `RewardConfig`. + """ + # Get reward class type + cls = reward_config["cls"] + if isinstance(cls, str): + obj = locate(cls) + assert isinstance(obj, type) and issubclass(obj, AbstractReward) + reward_config["cls"] = cls = obj + + # Get reward constructor keyword-arguments + kwargs = reward_config.get("kwargs", {}) + + # Special handling for `BaseMixtureReward` + if issubclass(cls, BaseMixtureReward): + for component_config in kwargs["components"]: + sanitize_reward_config(component_config) + + # Define helper to build the reward + def build_reward(env: InterfaceJiminyEnv, + reward_config: RewardConfig) -> AbstractReward: + """Instantiate a reward associated with a given environment provided + some reward configuration. + + :param env: Base environment or pipeline wrapper to wrap. + :param reward_config: Configuration of the reward, as a dict of type + `RewardConfig`. + """ + # Get reward class type + cls = reward_config["cls"] + assert isinstance(cls, type) and issubclass(cls, AbstractReward) + + # Get reward constructor keyword-arguments + kwargs = reward_config.get("kwargs", {}) + + # Special handling for `BaseMixtureReward` + if issubclass(cls, BaseMixtureReward): + kwargs["components"] = tuple( + build_reward(env, reward_config) + for reward_config in kwargs["components"]) + + # Special handling for `BaseQuantityReward` + if cls is BaseQuantityReward: + quantity_config = kwargs["quantity"] + kwargs["quantity"] = ( + quantity_config["cls"], quantity_config["kwargs"]) + + return cls(env, **kwargs) + + # Define helper to build reward + def build_composition(env_creator: Callable[..., InterfaceJiminyEnv], + reward_config: Optional[RewardConfig], + trajectories_config: Optional[TrajectoriesConfig], + **env_kwargs: Any) -> InterfaceJiminyEnv: + """Helper adding reward on top of a base environment or a pipeline + using `ComposedJiminyEnv` wrapper. + + :param env_creator: Callable that takes optional keyword arguments as + input and returns an pipeline or base environment. + :param reward_config: Configuration of the reward, as a dict of type + `RewardConfig`. + :param trajectories: Set of named trajectories as a dictionary. See + `ComposedJiminyEnv` documentation for details. + :param env_kwargs: Keyword arguments to forward to the constructor of + the wrapped environment. Note that it will only + overwrite the default value, so it will still be + possible to set different values by explicitly + defining them when calling the constructor of the + generated wrapper. + """ + # Instantiate the environment, which may be a lower-level wrapper + env = env_creator(**env_kwargs) + + # Instantiate the reward + reward = None + if reward_config is not None: + reward = build_reward(env, reward_config) + + # Get trajectory dataset + trajectories: Dict[str, Trajectory] = {} + if trajectories_config is not None: + trajectories = cast( + Dict[str, Trajectory], trajectories_config["dataset"]) + + # Instantiate the composition wrapper if necessary + if reward or trajectories: + env = ComposedJiminyEnv( + env, reward=reward, trajectories=trajectories) + + # Select the reference trajectory if specified + if trajectories_config is not None: + name = trajectories_config.get("name") + if name is not None: + mode = trajectories_config.get("mode", "raise") + env.quantities.select_trajectory(name, mode) + + return env + # Define helper to wrap a single layer def build_layer(env_creator: Callable[..., InterfaceJiminyEnv], wrapper_cls: Type[BasePipelineWrapper], @@ -120,15 +337,15 @@ def build_layer(env_creator: Callable[..., InterfaceJiminyEnv], :param env_creator: Callable that takes optional keyword arguments as input and returns an pipeline or base environment. + :param wrapper_cls: Type of wrapper to use to gather the environment + and the block. + :param wrapper_kwargs: Keyword arguments to forward to the constructor + of the wrapper. See 'env_kwargs'. :param block_cls: Type of block to connect to the environment, if any. `None` to disable. Optional: Disabled by default :param block_kwargs: Keyword arguments to forward to the constructor of the wrapped block. See 'env_kwargs'. - :param wrapper_cls: Type of wrapper to use to gather the environment - and the block. - :param wrapper_kwargs: Keyword arguments to forward to the constructor - of the wrapper. See 'env_kwargs'. :param env_kwargs: Keyword arguments to forward to the constructor of the wrapped environment. Note that it will only overwrite the default value, so it will still be @@ -170,15 +387,43 @@ def build_layer(env_creator: Callable[..., InterfaceJiminyEnv], # Instantiate the wrapper return wrapper_cls(*args, **wrapper_kwargs) - # Define callback for instantiating the base environment - env_cls: Union[Type[InterfaceJiminyEnv], str] = env_config["cls"] + # Define callable for instantiating the base environment + env_cls = env_config["cls"] if isinstance(env_cls, str): obj = locate(env_cls) - assert isinstance(obj, type) and issubclass(obj, InterfaceJiminyEnv) + assert isinstance(obj, type) and issubclass(obj, BaseJiminyEnv) env_cls = obj pipeline_creator: Callable[..., InterfaceJiminyEnv] = partial( env_cls, **env_config.get("kwargs", {})) + # Parse reward configuration + reward_config = env_config.get("reward") + if reward_config is not None: + sanitize_reward_config(reward_config) + + # Parse trajectory configuration + trajectories_config = env_config.get("trajectories") + if trajectories_config is not None: + trajectories = trajectories_config['dataset'] + assert isinstance(trajectories, dict) + for name, path_or_traj in trajectories.items(): + if isinstance(path_or_traj, Trajectory): + continue + path = pathlib.Path(path_or_traj) + if not path.is_absolute(): + if root_path is None: + raise RuntimeError( + "The argument 'root_path' must be provided when " + "specifying relative trajectory paths.") + path = pathlib.Path(root_path) / path + trajectories[name] = load_trajectory_from_hdf5(path) + + # Compose base environment with an extra user-specified reward + pipeline_creator = partial(build_composition, + pipeline_creator, + reward_config, + trajectories_config) + # Generate pipeline recursively for layer_config in layers_config: # Extract block and wrapper config @@ -239,15 +484,94 @@ def build_layer(env_creator: Callable[..., InterfaceJiminyEnv], return pipeline_creator -def load_pipeline(fullpath: str) -> Callable[..., InterfaceJiminyEnv]: +def load_pipeline(fullpath: Union[str, pathlib.Path] + ) -> Callable[..., InterfaceJiminyEnv]: """Load pipeline from JSON or TOML configuration file. :param: Fullpath of the configuration file. """ - file_ext = pathlib.Path(fullpath).suffix + fullpath = pathlib.Path(fullpath) + root_path, file_ext = fullpath.parent, fullpath.suffix with open(fullpath, 'r') as f: if file_ext == '.json': - return build_pipeline(**json.load(f)) + return build_pipeline(**json.load(f), root_path=root_path) if file_ext == '.toml': - return build_pipeline(**toml.load(f)) + return build_pipeline(**toml.load(f), root_path=root_path) raise ValueError("Only json and toml formats are supported.") + + +def save_trajectory_to_hdf5(trajectory: Trajectory, + fullpath: Union[str, pathlib.Path]) -> None: + """Export a trajectory object to HDF5 format. + + :param trajectory: Trajectory object to save. + :param fullpath: Fullpath of the generated HDF5 file. + """ + # Create HDF5 file + hdf_obj = h5py.File(fullpath, "w") + + # Dump each state attribute that are specified for all states at once + if trajectory.states: + state_dict = asdict(trajectory.states[0]) + state_fields = tuple( + key for key, value in state_dict.items() if value is not None) + for key in state_fields: + data = np.stack([ + getattr(state, key) for state in trajectory.states], axis=0) + hdf_obj.create_dataset(name=f"states/{key}", data=data) + + # Dump serialized robot + robot_data = jiminy.save_to_binary(trajectory.robot) + dataset = hdf_obj.create_dataset(name="robot", data=np.array(robot_data)) + + # Dump whether to use the theoretical model of the robot + dataset.attrs["use_theoretical_model"] = trajectory.use_theoretical_model + + # Close the HDF5 file + hdf_obj.close() + + +def load_trajectory_from_hdf5( + fullpath: Union[str, pathlib.Path]) -> Trajectory: + """Import a trajectory object from file in HDF5 format. + + :param fullpath: Fullpath of the HDF5 file to import. + + :returns: Loaded trajectory object. + """ + # Open HDF5 file + hdf_obj = h5py.File(fullpath, "r") + + # Get all state attributes that are specified + states_dict = {} + if 'states' in hdf_obj.keys(): + for key, value in hdf_obj['states'].items(): + states_dict[key] = value[...] + + # Re-construct state sequence + states = [] + for args in zip(*states_dict.values()): + states.append(State(**dict(zip(states_dict.keys(), args)))) + + # Build trajectory from data. + # Null char '\0' must be added at the end to match original string length. + dataset = hdf_obj['robot'] + robot_data = dataset[()] + robot_data += b'\0' * ( + dataset.nbytes - len(robot_data)) # pylint: disable=no-member + try: + robot = jiminy.load_from_binary(robot_data) + except MemoryError as e: + raise MemoryError( + "Impossible to build robot from serialized binary data. Make sure " + "that data has been generated on a machine with the same hardware " + "as this one.") from e + + # Load whether to use the theoretical model of the robot + use_theoretical_model = dataset.attrs["use_theoretical_model"] + + # Close the HDF5 file + hdf_obj.close() + + # Re-construct the whole trajectory + return Trajectory(states, robot, use_theoretical_model) diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py b/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py index 47d1958c2..dd8d61d43 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py @@ -408,7 +408,7 @@ def _reduce(fn_1, fn_2, delayed): fn_1(delayed) return partial(_reduce, fn_1, fn_2) if is_out_1 and not is_out_2: - if has_args: + if has_args: # pylint: disable=possibly-used-before-assignment def _reduce(fn_1, fn_2, field_2, args_2, delayed): fn_2(delayed[field_2], *args_2) fn_1(delayed) diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py index 1e70a55d5..1111b4863 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py @@ -1,16 +1,16 @@ # pylint: disable=missing-module-docstring from .observation_filter import FilterObservation -from .observation_stack import StackedJiminyEnv +from .observation_stack import StackObservation from .normalize import NormalizeAction, NormalizeObservation from .flatten import FlattenAction, FlattenObservation __all__ = [ 'FilterObservation', - 'StackedJiminyEnv', + 'StackObservation', 'NormalizeAction', 'NormalizeObservation', 'FlattenAction', - 'FlattenObservation', + 'FlattenObservation' ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py index 194f69f7f..e700a63b1 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py @@ -1,4 +1,6 @@ -""" TODO: Write documentation. +"""This module implements a wrapper for stacking observations over time. + +This wrapper """ import logging from collections import deque @@ -25,13 +27,18 @@ LOGGER = logging.getLogger(__name__) -class StackedJiminyEnv( +class StackObservation( BasePipelineWrapper[StackedObsT, ActT, NestedObsT, ActT], Generic[NestedObsT, ActT]): """Partially stack observations in a rolling manner. This wrapper combines and extends OpenAI Gym wrappers `FrameStack` and - `FilteredJiminyEnv` to support nested filter keys. + `FilteredJiminyEnv` to support nested filter keys. It derives from + `BasePipelineWrapper` rather than `gym.Wrapper`. This means that + observations can be stacked at observation update period rather than step + update period. It is also possible to access the stacked observation from + any block of the environment pipeline, and it will be taken into account + when calling `evaluate` or `play_interactive`. It adds one extra dimension to all the leaves of the original observation spaces that must be stacked. In such a case, the first dimension diff --git a/python/gym_jiminy/common/setup.py b/python/gym_jiminy/common/setup.py index 0f1dce609..0956bae00 100644 --- a/python/gym_jiminy/common/setup.py +++ b/python/gym_jiminy/common/setup.py @@ -49,8 +49,8 @@ package_data={"gym_jiminy.common": ["py.typed"]}, install_requires=[ f"jiminy-py~={jiminy_version}", - # Use to perform linear algebra computation - "numpy", + # Use internally for loading and dumping trajectories + "h5py", # Use internally to speedup computation of math methods # - 0.54: Adds 'np.clip' "numba>=0.54.0", diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py b/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py index da9e1a114..9cc85bf63 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py @@ -4,7 +4,6 @@ import numpy as np import gymnasium as gym -from gymnasium.spaces import flatten_space import jiminy_py.core as jiminy from jiminy_py.simulator import Simulator @@ -98,16 +97,15 @@ def __init__(self, robot.initialize( urdf_path, has_freeflyer=False, mesh_package_dirs=[data_dir]) - # Add motors and sensors + # Add motor motor_joint_name = "SecondArmJoint" - encoder_joint_names = ("FirstArmJoint", "SecondArmJoint") motor = jiminy.SimpleMotor(motor_joint_name) robot.attach_motor(motor) motor.initialize(motor_joint_name) - for joint_name in encoder_joint_names: - encoder = jiminy.EncoderSensor(joint_name) - robot.attach_sensor(encoder) - encoder.initialize(joint_name) + for joint_name in ("FirstArmJoint", "SecondArmJoint"): + sensor = jiminy.EncoderSensor(joint_name) + robot.attach_sensor(sensor) + sensor.initialize(joint_name=joint_name) # Instantiate simulator simulator = Simulator(robot, viewer_kwargs=viewer_kwargs) @@ -118,7 +116,7 @@ def __init__(self, # Map between discrete actions and actual motor torque if necessary if not self.continuous: - command_limit = np.asarray(motor.command_limit) + command_limit = np.array(motor.effort_limit) self.AVAIL_CTRL = (-command_limit, np.array(0.0), command_limit) # Internal parameters used for computing termination condition @@ -157,8 +155,14 @@ def _initialize_observation_space(self) -> None: Only the state is observable, while by default, the current time, state, and sensors data are available. """ - self.observation_space = flatten_space( - self._get_agent_state_space(use_theoretical_model=True)) + state_space = self._get_agent_state_space(use_theoretical_model=True) + position_space, velocity_space = state_space['q'], state_space['v'] + assert isinstance(position_space, gym.spaces.Box) + assert isinstance(velocity_space, gym.spaces.Box) + self.observation_space = gym.spaces.Box( + low=np.concatenate((position_space.low, velocity_space.low)), + high=np.concatenate((position_space.high, velocity_space.high)), + dtype=np.float64) def refresh_observation(self, measurement: EngineObsType) -> None: angles, velocities = measurement['measurements']['EncoderSensor'] @@ -192,17 +196,19 @@ def _sample_state(self) -> Tuple[np.ndarray, np.ndarray]: qvel = sample(scale=DTHETA_RANDOM_MAX, shape=(2,), rg=self.np_random) return qpos, qvel - def has_terminated(self) -> Tuple[bool, bool]: + def has_terminated(self, info: InfoType) -> Tuple[bool, bool]: """Determine whether the episode is over. It terminates (`terminated=True`) if the goal has been achieved, namely if the tip of the Acrobot is above 'HEIGHT_REL_DEFAULT_THRESHOLD'. Apart from that, there is no specific truncation condition. + :param info: Dictionary of extra information for monitoring. + :returns: terminated and truncated flags. """ # Call base implementation - terminated, truncated = super().has_terminated() + terminated, truncated = super().has_terminated(info) # Check if the agent has successfully solved the task tip_transform = self.robot.pinocchio_data.oMf[self._tipIndex] @@ -224,10 +230,7 @@ def compute_command(self, action: np.ndarray, command: np.ndarray) -> None: if ACTION_NOISE > 0.0: command += sample(scale=ACTION_NOISE, rg=self.np_random) - def compute_reward(self, - terminated: bool, - truncated: bool, - info: InfoType) -> float: + def compute_reward(self, terminated: bool, info: InfoType) -> float: """Compute reward at current episode state. Get a small negative reward till success. diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/ant.py b/python/gym_jiminy/envs/gym_jiminy/envs/ant.py index d7ed90c2a..d8084a412 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/ant.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/ant.py @@ -1,7 +1,13 @@ - -""" TODO: Write documentation. +"""This module implements to classic toy model environment "ant" using Jiminy +simulator for integrating the rigid-body dynamics. + +The agent is a basic 3D quadruped. This dynamics is very basic and not +realistic from a robotic standpoint. Its main advantage is that training a +policy for a given task is extremely rapidly compared to a real quadrupedal +robot such as Anymal. Still, it is complex enough for prototyping learning of +advanced motions or locomotion tasks, as well as model-based observers and +controllers usually intended for robotic applications. """ - import os import sys from typing import Any, Tuple, Sequence @@ -66,8 +72,7 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: simulator=simulator, debug=debug, **{**dict( - step_dt=STEP_DT, - enforce_bounded_spaces=False), + step_dt=STEP_DT), **kwargs}) # Define observation slices proxy for fast access. @@ -202,15 +207,19 @@ def refresh_observation(self, measurement: EngineObsType) -> None: self.observation_space.high, out=self.observation) - def has_terminated(self) -> Tuple[bool, bool]: + def has_terminated(self, info: InfoType) -> Tuple[bool, bool]: """Determine whether the episode is over. It adds one extra truncation criterion on top of the one defined in the base implementation. More precisely, the vertical height of the floating base of the robot must not exceed 0.2m. + + :param info: Dictionary of extra information for monitoring. + + :returns: terminated and truncated flags. """ # Call base implementation - terminated, truncated = super().has_terminated() + terminated, truncated = super().has_terminated(info) # Check if the agent is jumping far too high or stuck on its back zpos = self._robot_state_q[2] @@ -219,10 +228,7 @@ def has_terminated(self) -> Tuple[bool, bool]: return terminated, truncated - def compute_reward(self, - terminated: bool, - truncated: bool, - info: InfoType) -> float: + def compute_reward(self, terminated: bool, info: InfoType) -> float: """Compute reward at current episode state. The reward is defined as the sum of several individual components: diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py b/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py index dda5ef4eb..1865cc5d9 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/anymal.py @@ -16,12 +16,18 @@ # Default simulation duration (:float [s]) SIMULATION_DURATION = 20.0 + # Ratio between the High-level neural network PID target update and Low-level # PID torque update (:int [NA]) HLC_TO_LLC_RATIO = 1 + # Stepper update period (:float [s]) STEP_DT = 0.04 +# Motor safety to avoid violent motions +MOTOR_VELOCITY_MAX = 4.0 +MOTOR_ACCELERATION_MAX = 30.0 + # PID proportional gains (one per actuated joint) PD_KP = (1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0) @@ -89,9 +95,9 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: update_ratio=HLC_TO_LLC_RATIO, kp=PD_KP, kd=PD_KD, - target_position_margin=0.0, - target_velocity_limit=float("inf"), - target_acceleration_limit=float("inf") + joint_position_margin=0.0, + joint_velocity_limit=MOTOR_VELOCITY_MAX, + joint_acceleration_limit=MOTOR_ACCELERATION_MAX ) ), wrapper=dict( diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py b/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py index 99dfcd081..c50c89bfa 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py @@ -1,3 +1,5 @@ +""" TODO: Write documentation. +""" import os import sys from pathlib import Path @@ -5,8 +7,8 @@ import numpy as np -from jiminy_py.core import build_models_from_urdf, Robot -from jiminy_py.robot import load_hardware_description_file, BaseJiminyRobot +import jiminy_py.core as jiminy +from jiminy_py.robot import load_hardware_description_file from jiminy_py.viewer.viewer import DEFAULT_CAMERA_XYZRPY_REL from pinocchio import neutral, buildReducedModel @@ -33,9 +35,11 @@ # Stepper update period (:float [s]) STEP_DT = 0.04 +# Motor safety to avoid violent motions MOTOR_POSITION_MARGIN = 0.02 -MOTOR_VELOCITY_MAX = 3.0 -MOTOR_ACCELERATION_MAX = 40.0 +MOTOR_VELOCITY_SAFE_GAIN = 0.15 +MOTOR_VELOCITY_MAX = 4.0 +MOTOR_ACCELERATION_MAX = 30.0 # PID proportional gains (one per actuated joint) PD_REDUCED_KP = ( @@ -45,10 +49,11 @@ 5000.0, 5000.0, 8000.0, 4000.0, 8000.0, 5000.0) PD_REDUCED_KD = ( # Left leg: [HpX, HpZ, HpY, KnY, AkY, AkX] - 0.02, 0.01, 0.015, 0.01, 0.015, 0.01, + 0.02, 0.01, 0.02, 0.01, 0.025, 0.01, # Right leg: [HpX, HpZ, HpY, KnY, AkY, AkX] - 0.02, 0.01, 0.015, 0.01, 0.015, 0.01) + 0.02, 0.01, 0.02, 0.01, 0.025, 0.01) +# PID derivative gains (one per actuated joint) PD_FULL_KP = ( # Neck: [Y] 100.0, @@ -121,7 +126,7 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: """ # Get the urdf and mesh paths data_dir = str(files("gym_jiminy.envs") / "data/bipedal_robots/atlas") - urdf_path = os.path.join(data_dir, "atlas_v4.urdf") + urdf_path = os.path.join(data_dir, "atlas.urdf") # Override default camera pose to change the reference frame kwargs.setdefault("viewer_kwargs", {}).setdefault( @@ -168,14 +173,14 @@ class AtlasReducedJiminyEnv(WalkerJiminyEnv): def __init__(self, debug: bool = False, **kwargs: Any) -> None: # Get the urdf and mesh paths data_dir = str(files("gym_jiminy.envs") / "data/bipedal_robots/atlas") - urdf_path = os.path.join(data_dir, "atlas_v4.urdf") + urdf_path = os.path.join(data_dir, "atlas.urdf") # Load the full models - pinocchio_model, collision_model, visual_model = \ - build_models_from_urdf(urdf_path, - has_freeflyer=True, - build_visual_model=True, - mesh_package_dirs=[data_dir]) + pinocchio_model, collision_model, visual_model = ( + jiminy.build_models_from_urdf(urdf_path, + has_freeflyer=True, + build_visual_model=True, + mesh_package_dirs=[data_dir])) # Generate the reference configuration def joint_position_index(joint_name: str) -> int: @@ -206,9 +211,8 @@ def joint_position_index(joint_name: str) -> int: joint_locked_indices, qpos) # Build the robot and load the hardware - robot = BaseJiminyRobot() - Robot.initialize(robot, pinocchio_model, collision_model, visual_model) - robot._urdf_path_orig = urdf_path # type: ignore[attr-defined] + robot = jiminy.Robot() + robot.initialize(pinocchio_model, collision_model, visual_model) hardware_path = str(Path(urdf_path).with_suffix('')) + '_hardware.toml' load_hardware_description_file( robot, @@ -224,6 +228,8 @@ def joint_position_index(joint_name: str) -> int: avoid_instable_collisions=True, debug=debug, **{**dict( + config_path=str( + Path(urdf_path).with_suffix('')) + '_options.toml', simulation_duration_max=SIMULATION_DURATION, step_dt=STEP_DT, reward_mixture=REWARD_MIXTURE, @@ -244,7 +250,7 @@ def joint_position_index(joint_name: str) -> int: cls=MotorSafetyLimit, kwargs=dict( kp=1.0 / MOTOR_POSITION_MARGIN, - kd=1.0 / MOTOR_VELOCITY_MAX, + kd=MOTOR_VELOCITY_SAFE_GAIN, soft_position_margin=0.0, soft_velocity_max=MOTOR_VELOCITY_MAX, ) @@ -256,9 +262,9 @@ def joint_position_index(joint_name: str) -> int: update_ratio=1, kp=PD_FULL_KP, kd=PD_FULL_KD, - target_position_margin=0.0, - target_velocity_limit=MOTOR_VELOCITY_MAX, - target_acceleration_limit=MOTOR_ACCELERATION_MAX + joint_position_margin=0.0, + joint_velocity_limit=MOTOR_VELOCITY_MAX, + joint_acceleration_limit=MOTOR_ACCELERATION_MAX ) ), wrapper=dict( @@ -314,9 +320,9 @@ def joint_position_index(joint_name: str) -> int: update_ratio=1, kp=PD_REDUCED_KP, kd=PD_REDUCED_KD, - target_position_margin=0.0, - target_velocity_limit=MOTOR_VELOCITY_MAX, - target_acceleration_limit=MOTOR_ACCELERATION_MAX + joint_position_margin=0.0, + joint_velocity_limit=MOTOR_VELOCITY_MAX, + joint_acceleration_limit=MOTOR_ACCELERATION_MAX ) ), wrapper=dict( diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py b/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py index af8195beb..5ea599d06 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py @@ -116,29 +116,26 @@ def __init__(self, # Add motors and sensors motor_joint_name = "slider_to_cart" - encoder_sensors_descr = { - "slider": "slider_to_cart", - "pole": "cart_to_pole" - } motor = jiminy.SimpleMotor(motor_joint_name) robot.attach_motor(motor) motor.initialize(motor_joint_name) - for sensor_name, joint_name in encoder_sensors_descr.items(): + for sensor_name, joint_name in ( + ("slider", "slider_to_cart"), ("pole", "cart_to_pole")): encoder = jiminy.EncoderSensor(sensor_name) robot.attach_sensor(encoder) - encoder.initialize(joint_name) + encoder.initialize(joint_name=joint_name) # Instantiate simulator simulator = Simulator(robot, viewer_kwargs=viewer_kwargs) # OpenAI Gym implementation of Cartpole has no velocity limit - model_options = simulator.robot.get_model_options() - model_options["joints"]["enableVelocityLimit"] = False - simulator.robot.set_model_options(model_options) + motor_options = motor.get_options() + motor_options["enableVelocityLimit"] = False + motor.set_options(motor_options) # Map between discrete actions and actual motor force if necessary if not self.continuous: - command_limit = np.asarray(motor.command_limit) + command_limit = np.array(motor.effort_limit) self.AVAIL_CTRL = (-command_limit, np.array(0.0), command_limit) # Configure the learning environment @@ -163,14 +160,18 @@ def _initialize_observation_space(self) -> None: See documentation: https://gym.openai.com/envs/CartPole-v1/. """ - # Compute observation bounds - high = np.array([X_THRESHOLD, - THETA_THRESHOLD, - *self.robot.velocity_limit]) + # Define custom symmetric position bounds + position_limit_upper = np.array([X_THRESHOLD, THETA_THRESHOLD]) + + # Get velocity bounds associated the theoretical model + velocity_limit = self.robot.get_theoretical_velocity_from_extended( + self.robot.pinocchio_model.velocityLimit) # Set the observation space + state_limit_upper = np.concatenate(( + position_limit_upper, velocity_limit)) self.observation_space = spaces.Box( - low=-high, high=high, dtype=np.float64) + low=-state_limit_upper, high=state_limit_upper, dtype=np.float64) def _initialize_action_space(self) -> None: """ TODO: Write documentation. @@ -188,8 +189,9 @@ def _sample_state(self) -> Tuple[np.ndarray, np.ndarray]: Bounds of hypercube associated with initial state of robot. """ - qpos = sample(scale=np.array([ - X_RANDOM_MAX, THETA_RANDOM_MAX]), rg=self.np_random) + x, theta = sample(scale=np.array( + [X_RANDOM_MAX, THETA_RANDOM_MAX]), rg=self.np_random) + qpos = np.array([x, np.cos(theta), np.sin(theta)]) qvel = sample(scale=np.array([ DX_RANDOM_MAX, DTHETA_RANDOM_MAX]), rg=self.np_random) return qpos, qvel @@ -207,10 +209,7 @@ def compute_command(self, action: np.ndarray, command: np.ndarray) -> None: """ command[:] = action if self.continuous else self.AVAIL_CTRL[action] - def compute_reward(self, - terminated: bool, - truncated: bool, - info: InfoType) -> float: + def compute_reward(self, terminated: bool, info: InfoType) -> float: """Compute reward at current episode state. Constant positive reward equal to 1.0 as long as no termination diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/cassie.py b/python/gym_jiminy/envs/gym_jiminy/envs/cassie.py index b13c5572c..01df9747f 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/cassie.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/cassie.py @@ -1,3 +1,5 @@ +""" TODO: Write documentation. +""" import os import sys import math @@ -6,12 +8,8 @@ import numpy as np -from jiminy_py.core import (JointModelType, - get_joint_type, - build_models_from_urdf, - DistanceConstraint, - Robot) -from jiminy_py.robot import load_hardware_description_file, BaseJiminyRobot +import jiminy_py.core as jiminy +from jiminy_py.robot import load_hardware_description_file from pinocchio import neutral, SE3, buildReducedModel from gym_jiminy.common.envs import WalkerJiminyEnv @@ -40,9 +38,14 @@ # Stepper update period (:float [s]) STEP_DT = 0.04 +# Motor safety to avoid violent motions +MOTOR_VELOCITY_MAX = 4.0 +MOTOR_ACCELERATION_MAX = 30.0 + # PID proportional gains (one per actuated joint) -PD_KP = (100.0, 100.0, 100.0, 100.0, 80.0, - 100.0, 100.0, 100.0, 100.0, 80.0) +PD_KP = (4.0, 4.0, 6.25, 6.25, 1.6, + 4.0, 4.0, 6.25, 6.25, 1.6) + # PID derivative gains (one per actuated joint) PD_KD = (0.02, 0.02, 0.02, 0.02, 0.015, 0.02, 0.02, 0.02, 0.02, 0.015) @@ -81,11 +84,11 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: urdf_path = os.path.join(data_dir, "cassie.urdf") # Load the full models - pinocchio_model, collision_model, visual_model = \ - build_models_from_urdf(urdf_path, - has_freeflyer=True, - build_visual_model=True, - mesh_package_dirs=[data_dir]) + pinocchio_model, collision_model, visual_model = ( + jiminy.build_models_from_urdf(urdf_path, + has_freeflyer=True, + build_visual_model=True, + mesh_package_dirs=[data_dir])) # Fix passive rotary joints with spring. # Alternatively, it would be more realistic to model them using the @@ -100,9 +103,8 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: joint_locked_indices, qpos) # Build the robot and load the hardware - robot = BaseJiminyRobot() - Robot.initialize(robot, pinocchio_model, collision_model, visual_model) - robot._urdf_path_orig = urdf_path # type: ignore[attr-defined] + robot = jiminy.Robot() + robot.initialize(pinocchio_model, collision_model, visual_model) hardware_path = str(Path(urdf_path).with_suffix('')) + '_hardware.toml' load_hardware_description_file( robot, @@ -118,6 +120,8 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: avoid_instable_collisions=True, debug=debug, **{**dict( + config_path=str( + Path(urdf_path).with_suffix('')) + '_options.toml', simulation_duration_max=SIMULATION_DURATION, step_dt=STEP_DT, reward_mixture=REWARD_MIXTURE, @@ -133,7 +137,7 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: "right_pushrod_tarsus", "right_tarsus", M_pushrod_tarsus_right) self.robot.add_frame( "right_pushrod_hip", "hip_flexion_right", M_pushrod_hip_right) - pushrod_right = DistanceConstraint( + pushrod_right = jiminy.DistanceConstraint( "right_pushrod_tarsus", "right_pushrod_hip") pushrod_right.baumgarte_freq = 20.0 self.robot.add_constraint("pushrod_right", pushrod_right) @@ -145,7 +149,7 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: "left_pushrod_tarsus", "left_tarsus", M_pushrod_tarsus_left) self.robot.add_frame( "left_pushrod_hip", "hip_flexion_left", M_pushrod_hip_left) - pushrod_left = DistanceConstraint( + pushrod_left = jiminy.DistanceConstraint( "left_pushrod_tarsus", "left_pushrod_hip") pushrod_left.baumgarte_freq = 20.0 self.robot.add_constraint("pushrod_left", pushrod_left) @@ -161,10 +165,10 @@ def set_joint_rotary_position(joint_name: str, theta: float) -> None: """ joint_index = self.robot.pinocchio_model.getJointId(joint_name) joint = self.robot.pinocchio_model.joints[joint_index] - joint_type = get_joint_type( + joint_type = jiminy.get_joint_type( self.robot.pinocchio_model, joint_index) q_joint: Union[Sequence[float], float] - if joint_type == JointModelType.ROTARY_UNBOUNDED: + if joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: q_joint = (math.cos(theta), math.sin(theta)) else: q_joint = theta @@ -195,9 +199,9 @@ def set_joint_rotary_position(joint_name: str, theta: float) -> None: update_ratio=HLC_TO_LLC_RATIO, kp=PD_KP, kd=PD_KD, - target_position_margin=0.0, - target_velocity_limit=float("inf"), - target_acceleration_limit=float("inf") + joint_position_margin=0.0, + joint_velocity_limit=MOTOR_VELOCITY_MAX, + joint_acceleration_limit=MOTOR_ACCELERATION_MAX ) ), wrapper=dict( diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/digit.py b/python/gym_jiminy/envs/gym_jiminy/envs/digit.py index 3a6c2895c..314f4d093 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/digit.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/digit.py @@ -1,3 +1,5 @@ +""" TODO: Write documentation. +""" import os import sys import math @@ -40,16 +42,21 @@ # Stepper update period (:float [s]) STEP_DT = 0.04 +# Motor safety to avoid violent motions +MOTOR_VELOCITY_MAX = 4.0 +MOTOR_ACCELERATION_MAX = 30.0 + # PID proportional gains (one per actuated joint) PD_KP = ( # Left leg: [HpX, HpZ, HpY, KnY, AkY, AkX] - 80.0, 100.0, 100.0, 100.0, 80.0, 80.0, + 1.0, 2.0, 6.25, 6.25, 1.6, 1.6, # Left arm: [ShX, ShY, ShZ, ElY] - 50.0, 50.0, 50.0, 50.0, + 0.625, 0.625, 1.0, 0.625, # Right leg: [HpX, HpZ, HpY, KnY, AkY, AkX] - 80.0, 100.0, 100.0, 100.0, 80.0, 80.0, + 1.0, 2.0, 6.25, 6.25, 1.6, 1.6, # Right arm: [ShX, ShY, ShZ, ElY] - 50.0, 50.0, 50.0, 50.0) + 0.625, 0.625, 1.0, 0.625) + # PID derivative gains (one per actuated joint) PD_KD = ( # Left leg: [HpX, HpZ, HpY, KnY, AkY, AkX] @@ -210,9 +217,9 @@ def set_joint_rotary_position(joint_name: str, theta: float) -> None: update_ratio=HLC_TO_LLC_RATIO, kp=PD_KP, kd=PD_KD, - target_position_margin=0.0, - target_velocity_limit=float("inf"), - target_acceleration_limit=float("inf") + joint_position_margin=0.0, + joint_velocity_limit=MOTOR_VELOCITY_MAX, + joint_acceleration_limit=MOTOR_ACCELERATION_MAX ) ), wrapper=dict( diff --git a/python/gym_jiminy/examples/quantity_benchmark.py b/python/gym_jiminy/examples/quantity_benchmark.py index 5c381d080..4e6a46020 100644 --- a/python/gym_jiminy/examples/quantity_benchmark.py +++ b/python/gym_jiminy/examples/quantity_benchmark.py @@ -4,14 +4,14 @@ import matplotlib.pyplot as plt import gymnasium as gym -import gym_jiminy.common.bases.quantity -from gym_jiminy.common.quantities import QuantityManager, EulerAnglesFrame +import gym_jiminy.common.bases.quantities +from gym_jiminy.common.quantities import QuantityManager, FrameEulerAngles # Define number of samples for benchmarking N_SAMPLES = 50000 # Disable caching by forcing `SharedCache.has_value` to always return `False` -setattr(gym_jiminy.common.bases.quantity.SharedCache, +setattr(gym_jiminy.common.bases.quantities.SharedCache, "has_value", property(lambda self: False)) @@ -25,7 +25,7 @@ quantity_manager = QuantityManager( env.simulator, { - f"rpy_{i}": (EulerAnglesFrame, dict(frame_name=frame.name)) + f"rpy_{i}": (FrameEulerAngles, dict(frame_name=frame.name)) for i, frame in enumerate(env.robot.pinocchio_model.frames) }) @@ -36,12 +36,12 @@ quantity_manager.reset(reset_tracking=True) # Fetch all quantities once to update dynamic computation graph - for j, quantity in enumerate(quantity_manager._quantities.values()): + for j, quantity in enumerate(quantity_manager.registry.values()): quantity.get() if i == j + 1: break - # Extract batched data buffer of `EulerAnglesFrame` quantities + # Extract batched data buffer of `FrameEulerAngles` quantities shared_data = quantity.requirements['data'] # Benchmark computation of batched data buffer diff --git a/python/gym_jiminy/examples/stable_baselines3/tools/utilities.py b/python/gym_jiminy/examples/stable_baselines3/tools/utilities.py index 753756220..9b34536e8 100644 --- a/python/gym_jiminy/examples/stable_baselines3/tools/utilities.py +++ b/python/gym_jiminy/examples/stable_baselines3/tools/utilities.py @@ -133,7 +133,7 @@ def test(test_agent: BaseAlgorithm, Optional: True by default. """ # Check user arguments - if (math.isinf(max_episodes) and math.isinf(max_duration)): + if math.isinf(max_episodes) and math.isinf(max_duration): raise ValueError( "Either 'max_episodes' or 'max_duration' must be finite.") diff --git a/python/gym_jiminy/examples/tutorial_rl.ipynb b/python/gym_jiminy/examples/tutorial_rl.ipynb index 94de44937..6dad16b4b 100644 --- a/python/gym_jiminy/examples/tutorial_rl.ipynb +++ b/python/gym_jiminy/examples/tutorial_rl.ipynb @@ -109,12 +109,12 @@ " qvel = sample(scale=np.array([0.05, 0.05]), rg=self.np_random)\n", " return qpos, qvel\n", "\n", - " def has_terminated(self):\n", + " def has_terminated(self, info):\n", " terminated = abs(self.robot_state.q[0]) > 4.0\n", " truncated = self.stepper_state.t > 30.0\n", " return terminated, truncated\n", "\n", - " def compute_reward(self, terminated, truncated, info):\n", + " def compute_reward(self, terminated, info):\n", " theta = self.robot_state.q[1]\n", " error = min(theta, 2 * math.pi - abs(theta))\n", " alpha = 1.0\n", diff --git a/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py b/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py index e25488060..bed833d9f 100644 --- a/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py +++ b/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py @@ -29,6 +29,7 @@ from ray.rllib.algorithms.ppo.ppo_torch_policy import ( PPOTorchPolicy as _PPOTorchPolicy) from ray.rllib.algorithms.ppo.ppo_tf_policy import validate_config +from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.utils.numpy import convert_to_numpy from ray.rllib.utils.torch_utils import l2_loss from ray.rllib.utils.annotations import override @@ -202,6 +203,7 @@ def __init__(self, algo_class: Optional[Type["PPO"]] = None): self.temporal_barrier_threshold = float('inf') self.temporal_barrier_reg = 0.0 self.symmetric_policy_reg = 0.0 + self.enable_symmetry_surrogate_loss = False self.caps_temporal_reg = 0.0 self.caps_spatial_reg = 0.0 self.caps_global_reg = 0.0 @@ -219,6 +221,7 @@ def training( temporal_barrier_threshold: Optional[float] = None, temporal_barrier_reg: Optional[float] = None, symmetric_policy_reg: Optional[float] = None, + enable_symmetry_surrogate_loss: Optional[bool] = None, caps_temporal_reg: Optional[float] = None, caps_spatial_reg: Optional[float] = None, caps_global_reg: Optional[float] = None, @@ -244,6 +247,9 @@ def training( self.temporal_barrier_reg = temporal_barrier_reg if symmetric_policy_reg is not None: self.symmetric_policy_reg = symmetric_policy_reg + if enable_symmetry_surrogate_loss is not None: + self.enable_symmetry_surrogate_loss = \ + enable_symmetry_surrogate_loss if caps_temporal_reg is not None: self.caps_temporal_reg = caps_temporal_reg if caps_spatial_reg is not None: @@ -285,14 +291,35 @@ class PPOTorchPolicy(_PPOTorchPolicy): More specifically, it adds: - CAPS regularization, which combines the spatial and temporal - difference between previous and current state + difference between previous and current state. - Global regularization, which is the average norm of the action - temporal barrier, which is exponential barrier loss when the normalized action is above a threshold (much like interior point methods). - symmetry regularization, which is the error between actions and symmetric actions associated with symmetric observations. - - L2 regularization of policy network weights + - symmetry surrogate loss, which is the surrogate loss associated + with the symmetric (actions, observations) spaces. As the surrogate + loss goal is to increase the likelihood of selecting higher reward + actions given the current state, the symmetry surrogate loss enables + equivalent likelihood increase for selecting the symmetric higher + reward actions given the symmetric state. + - L2 regularization of policy network weights. + + More insights on the regularization losses with their emerging properties, + and on how to tune the parameters can be found in the reference articles: + - A. Duburcq, F. Schramm, G. Boeris, N. Bredeche, and Y. Chevaleyre, + “Reactive Stepping for Humanoid Robots using Reinforcement Learning: + Application to Standing Push Recovery on the Exoskeleton Atalante,” in + International Conference on Intelligent Robots and Systems (IROS), + vol. 2022-Octob. IEEE, oct 2022, pp. 9302–9309 + - S. Mysore, B. Mabsout, R. Mancuso, and K. Saenko, “Regularizing + action policies for smooth control with reinforcement learning,” + IEEE International Conference on Robotics and Automation (ICRA), + pp. 1810–1816, 2021 + - M. Mittal, N. Rudin, V. Klemm, A. Allshire, and M. Hutter, + “Symmetry considerations for learning task symmetric robot policies,” + arXiv preprint arXiv:2403.04359, 2024. """ def __init__(self, observation_space: gym.spaces.Space, @@ -335,7 +362,8 @@ def __init__(self, Dict[str, torch.Tensor], torch.Tensor]] = None self.action_mirror_mat: Optional[Union[ Dict[str, torch.Tensor], torch.Tensor]] = None - if config_dict["symmetric_policy_reg"] > 0.0: + if config_dict["symmetric_policy_reg"] > 0.0 or \ + config_dict["enable_symmetry_surrogate_loss"]: # Observation space is_obs_dict = hasattr(observation_space, "original_space") if is_obs_dict: @@ -386,6 +414,8 @@ def loss(self, train_batch: SampleBatch) -> Union[TensorType, List[TensorType]]: """Compute PPO loss with additional regularizations. """ + # pylint: disable=possibly-used-before-assignment + with torch.no_grad(): # Extract some proxies from convenience observation_true = train_batch["obs"] @@ -435,11 +465,12 @@ def loss(self, # Append the training batches to the set train_batches["noisy"] = train_batch_copy - if self.config["symmetric_policy_reg"] > 0.0: + if self.config["symmetric_policy_reg"] > 0.0 or \ + self.config["enable_symmetry_surrogate_loss"]: # Shallow copy the original training batch train_batch_copy = train_batch.copy(shallow=True) - # Compute mirrorred observation + # Compute mirrored observation assert self.obs_mirror_mat is not None assert isinstance(self.observation_space, gym.spaces.Box) observation_mirror = _compute_mirrored_value( @@ -521,7 +552,8 @@ def value_function( model, dist_class, action_worst_logits) # Compute the mirrored mean action corresponding to the mirrored action - if self.config["symmetric_policy_reg"] > 0.0: + if self.config["symmetric_policy_reg"] > 0.0 or \ + self.config["enable_symmetry_surrogate_loss"]: assert self.action_mirror_mat is not None assert isinstance(self.action_space, gym.spaces.Box) action_mirror_logits = action_logits["mirrored"] @@ -595,6 +627,36 @@ def value_function( total_loss += \ self.config["symmetric_policy_reg"] * symmetric_policy_reg + if self.config["enable_symmetry_surrogate_loss"]: + assert self.action_mirror_mat is not None + assert isinstance(self.action_space, gym.spaces.Box) + + # Get the mirror policy probability distribution + # i.e. ( x -> pi( x | mirrored observation ) ) + curr_action_mirror_dist = dist_class(action_mirror_logits, model) + + # The implementation assumes, at any time t, under the policy, + # the probability to be in state_t is equal to the probability to + # be in the mirrored state_t. Otherwise, their ratio needs to be + # added. + mirror_logp_ratio = torch.exp( + curr_action_mirror_dist.logp( + _compute_mirrored_value(train_batch[SampleBatch.ACTIONS], + self.action_space, + self.action_mirror_mat)) - + train_batch[SampleBatch.ACTION_LOGP]) + + symmetry_surrogate_loss = torch.mean(torch.min( + train_batch[Postprocessing.ADVANTAGES] * mirror_logp_ratio, + train_batch[Postprocessing.ADVANTAGES] * torch.clamp( + mirror_logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"]))) + + # Add symmetry surrogate loss to total loss + stats["symmetry_surrogate_loss"] = symmetry_surrogate_loss + total_loss -= symmetry_surrogate_loss + if self.config["l2_reg"] > 0.0: # Add actor l2-regularization loss l2_reg = 0.0 @@ -615,6 +677,9 @@ def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: """ stats_dict = super().stats_fn(train_batch) + if self.config["enable_symmetry_surrogate_loss"]: + stats_dict["symmetry_surrogate_loss"] = torch.mean( + torch.stack(self.get_tower_stats("symmetry_surrogate_loss"))) if self.config["symmetric_policy_reg"] > 0.0: stats_dict["symmetry"] = torch.mean( torch.stack(self.get_tower_stats("symmetric_policy_reg"))) diff --git a/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py b/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py index 40b806b93..fdd844d50 100644 --- a/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py +++ b/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py @@ -1,4 +1,10 @@ -""" TODO: Write documentation. +"""This module provides helper methods to make it easier to run learning +pipelines using Ray-RLlib framework. + +The main focus are: + * making it easy to build a policy from a checkpoint + * enabling evaluating the performance of a given policy without having to + fire up a whole ray server instance, which is normally required """ import os import re @@ -487,7 +493,8 @@ def train(algo: Algorithm, # Backup the policy if checkpoint_period > 0 and iter_num % checkpoint_period == 0: - algo.save() + algo.save( + checkpoint_dir=f"{algo.logdir}/checkpoint_{iter_num:06d}") # Check terminal conditions if 0 < max_timesteps < result["timesteps_total"]: @@ -510,7 +517,8 @@ def train(algo: Algorithm, result_logger.close() # Backup trained agent and return file location - return algo.save() + return algo.save( + checkpoint_dir=f"{algo.logdir}/checkpoint_{iter_num:06d}") def _restore_default_connectors() -> None: diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py index 91c308d15..828c16a67 100644 --- a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py @@ -1,4 +1,5 @@ -""" TODO: Write documentation. +"""This module implements a wrapper for maintaining a stable framerate when +human rendering is enabled during an episode. """ import time from typing import Any, List, Optional, Tuple, Generic, Union, SupportsFloat @@ -21,8 +22,8 @@ class FrameRateLimiter(gym.Wrapper, # [ObsT, ActT, ObsT, ActT], This wrapper only affects `render`, not `replay`. .. warning:: - This wrapper is only compatible with `BasePipelineWrapper` and - `BaseJiminyEnv` as it requires having a `step_dt` attribute. + This wrapper is only compatible with `InterfaceJiminyEnv` because it + requires having access to `step_dt` attribute. """ def __init__(self, # pylint: disable=unused-argument env: InterfaceJiminyEnv[ObsT, ActT], diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/meta_envs.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/meta_envs.py index 58e1a548f..74093d93f 100644 --- a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/meta_envs.py +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/meta_envs.py @@ -61,6 +61,8 @@ def __init__(self, ) -> None: """ TODO: Write documentation. """ + # pylint: disable=possibly-used-before-assignment + # Call base implementation super().__init__(env) diff --git a/python/gym_jiminy/unit_py/data/anymal_pipeline.json b/python/gym_jiminy/unit_py/data/anymal_pipeline.json index ea19558d5..6ca97568d 100644 --- a/python/gym_jiminy/unit_py/data/anymal_pipeline.json +++ b/python/gym_jiminy/unit_py/data/anymal_pipeline.json @@ -13,9 +13,9 @@ "update_ratio": 2, "kp": [1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0], "kd": [0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01], - "target_position_margin": 0.0, - "target_velocity_limit": 100.0, - "target_acceleration_limit": 10000.0 + "joint_position_margin": 0.0, + "joint_velocity_limit": 100.0, + "joint_acceleration_limit": 10000.0 } }, "wrapper": { @@ -51,7 +51,7 @@ }, { "wrapper": { - "cls": "gym_jiminy.common.wrappers.StackedJiminyEnv", + "cls": "gym_jiminy.common.wrappers.StackObservation", "kwargs": { "nested_filter_keys": [ ["t"], diff --git a/python/gym_jiminy/unit_py/data/anymal_pipeline.toml b/python/gym_jiminy/unit_py/data/anymal_pipeline.toml index ed567c463..4ff044484 100644 --- a/python/gym_jiminy/unit_py/data/anymal_pipeline.toml +++ b/python/gym_jiminy/unit_py/data/anymal_pipeline.toml @@ -2,6 +2,21 @@ cls = "gym_jiminy.envs.ANYmalJiminyEnv" [env_config.kwargs] step_dt = 0.04 +[env_config.reward] +cls = "gym_jiminy.common.compositions.AdditiveMixtureReward" +[env_config.trajectories] +mode = "raise" +name = "reference" +dataset.reference = "./anymal_trajectory.hdf5" +[env_config.reward.kwargs] +name = "reward_total" +weights = [0.6, 0.4] +[[env_config.reward.kwargs.components]] +cls = "gym_jiminy.common.compositions.OdometryVelocityReward" +[env_config.reward.kwargs.components.kwargs] +cutoff = 0.5 +[[env_config.reward.kwargs.components]] +cls = "gym_jiminy.common.compositions.SurviveReward" [[layers_config]] block.cls = "gym_jiminy.common.blocks.PDController" @@ -9,9 +24,9 @@ block.cls = "gym_jiminy.common.blocks.PDController" update_ratio = 2 kp = [1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0, 1500.0] kd = [0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01] -target_position_margin = 0.0 -target_velocity_limit = 100.0 -target_acceleration_limit = 10000.0 +joint_position_margin = 0.0 +joint_velocity_limit = 100.0 +joint_acceleration_limit = 10000.0 [layers_config.wrapper.kwargs] augment_observation = true @@ -32,7 +47,7 @@ kp = 1.0 ki = 0.1 [[layers_config]] -wrapper.cls = "gym_jiminy.common.wrappers.StackedJiminyEnv" +wrapper.cls = "gym_jiminy.common.wrappers.StackObservation" [layers_config.wrapper.kwargs] nested_filter_keys = [ ["t"], diff --git a/python/gym_jiminy/unit_py/test_deformation_estimator.py b/python/gym_jiminy/unit_py/test_deformation_estimator.py index 924e89777..3ae719563 100644 --- a/python/gym_jiminy/unit_py/test_deformation_estimator.py +++ b/python/gym_jiminy/unit_py/test_deformation_estimator.py @@ -26,8 +26,8 @@ class DeformationEstimatorBlock(unittest.TestCase): def _test_deformation_estimate(self, env, imu_atol, flex_atol): # Check that quaternion estimates from MahonyFilter are valid true_imu_rots = [] - for frame_name in env.robot.sensor_names['ImuSensor']: - frame_index = env.robot.pinocchio_model.getFrameId(frame_name) + for imu_sensor in env.robot.sensors['ImuSensor']: + frame_index = imu_sensor.frame_index frame_rot = env.robot.pinocchio_data.oMf[frame_index].rotation true_imu_rots.append(frame_rot) true_imu_quats = matrices_to_quat(tuple(true_imu_rots)) @@ -69,6 +69,9 @@ def test_arm(self): # Add motor motor_joint_name = 'base_to_link1' motor = jiminy.SimpleMotor(motor_joint_name) + motor_options = motor.get_options() + motor_options["enableVelocityLimit"] = False + motor.set_options(motor_options) robot.attach_motor(motor) motor.initialize(motor_joint_name) @@ -92,7 +95,6 @@ def test_arm(self): 'inertia': np.array([1.0, 1.0, 0.0]) } for i in range(1, 5) ] - model_options['joints']['enableVelocityLimit'] = False robot.set_model_options(model_options) # Create a simulator using this robot and controller @@ -136,7 +138,8 @@ def test_arm(self): deformation_estimator = DeformationEstimator( "deformation_estimator", env, - imu_frame_names=robot.sensor_names['ImuSensor'], + imu_frame_names=tuple( + sensor.name for sensor in robot.sensors['ImuSensor']), flex_frame_names=robot.flexibility_joint_names, ignore_twist=True, update_ratio=-1) @@ -219,7 +222,7 @@ def _setup(self) -> None: block=dict( cls=PDController, kwargs=dict( - kp=150.0, + kp=1.0, kd=0.03, update_ratio=1, ) diff --git a/python/gym_jiminy/unit_py/test_misc.py b/python/gym_jiminy/unit_py/test_misc.py new file mode 100644 index 000000000..ed341364b --- /dev/null +++ b/python/gym_jiminy/unit_py/test_misc.py @@ -0,0 +1,90 @@ +""" TODO: Write documentation +""" +import os +import sys +import shutil +import tempfile +import unittest + +import numpy as np +import gymnasium as gym + +from jiminy_py.core import EncoderSensor, ImuSensor, EffortSensor, ForceSensor +from jiminy_py.simulator import Simulator +from jiminy_py.viewer import Viewer +from jiminy_py.viewer.replay import ( + extract_replay_data_from_log, play_trajectories) + +if sys.version_info < (3, 9): + from importlib_resources import files +else: + from importlib.resources import files + + +class Miscellaneous(unittest.TestCase): + def test_play_log_data(self): + """ TODO: Write documentation. + """ + # Instantiate the environment + env = gym.make("gym_jiminy.envs:atlas", debug=True) + + # Run a few steps + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + env.stop() + + # Generate temporary video file + fd, video_path = tempfile.mkstemp(prefix=f"atlas_", suffix=".mp4") + os.close(fd) + + # Play trajectory + Viewer.close() + trajectory, update_hook, kwargs = extract_replay_data_from_log( + env.log_data) + viewer, = play_trajectories(trajectory, + update_hook, + backend="panda3d-sync", + record_video_path=video_path, + display_contacts=True, + display_f_external=True) + Viewer.close() + + # Check the external forces has been updated + np.testing.assert_allclose( + tuple(f_ext.vector for f_ext in env.robot_state.f_external[1:]), + tuple(f_ext.vector for f_ext in viewer.f_external)) + + # Check that sensor data has been updated + for key, data in env.robot.sensor_measurements.items(): + np.testing.assert_allclose( + data, trajectory.robot.sensor_measurements[key]) + + def test_default_hardware(self): + with tempfile.TemporaryDirectory() as tmpdirname: + # Copy original Atlas URDF and meshes + data_dir = files("gym_jiminy.envs") / "data/bipedal_robots/atlas" + urdf_path = os.path.join(tmpdirname, "atlas.urdf") + shutil.copy2(str(data_dir / "atlas.urdf"), urdf_path) + shutil.copytree( + str(data_dir / "meshes"), os.path.join(tmpdirname, "meshes")) + + # Build simulator with default hardware + simulator = Simulator.build(urdf_path) + + # Check that all mechanical joints are actuated + assert simulator.robot.nmotors == len( + simulator.robot.mechanical_joint_names) + + # Check that all mechanical joints have encoder and effort sensors + assert len(simulator.robot.sensors[EncoderSensor.type]) == 30 + assert len(simulator.robot.sensors[EffortSensor.type]) == 30 + + # Check that the root joint has an IMU sensor + assert len(simulator.robot.sensors[ImuSensor.type]) == 1 + sensor = simulator.robot.sensors[ImuSensor.type][0] + assert sensor.frame_name == ( + simulator.robot.pinocchio_model.frames[2].name) + + # Check that the each foot has a force sensor + assert len(simulator.robot.sensors[ForceSensor.type]) == 2 diff --git a/python/gym_jiminy/unit_py/test_pipeline_control.py b/python/gym_jiminy/unit_py/test_pipeline_control.py index d2bad793d..c34202b7d 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_control.py +++ b/python/gym_jiminy/unit_py/test_pipeline_control.py @@ -13,7 +13,7 @@ import matplotlib.pyplot as plt from PIL import Image -from jiminy_py.core import ImuSensor as imu +from jiminy_py.core import ImuSensor from jiminy_py.viewer import Viewer from gym_jiminy.envs import ( @@ -47,7 +47,7 @@ def _test_pid_standing(self): action = np.zeros(self.env.robot.nmotors) # Run the simulation - while self.env.stepper_state.t < 19.0: + while self.env.stepper_state.t < 9.0: self.env.step(action) # Export figure @@ -141,7 +141,7 @@ def test_mahony_filter(self): action[robot.get_motor(name).index] = value # Extract proxies for convenience - sensor = robot.get_sensor(imu.type, robot.sensor_names[imu.type][0]) + sensor = next(iter(robot.sensors[ImuSensor.type])) imu_rot = robot.pinocchio_data.oMf[sensor.frame_index].rotation # Check that the estimate IMU orientation is accurate over the episode @@ -283,17 +283,17 @@ def test_pd_controller(self): np.testing.assert_allclose(target_vel[(update_ratio-1)::update_ratio], command_vel[(update_ratio-1)::update_ratio], atol=TOLERANCE) - np.testing.assert_allclose(target_accel_diff, target_accel[1:], atol=TOLERANCE) - np.testing.assert_allclose(target_vel_diff, target_vel[1:], atol=TOLERANCE) + np.testing.assert_allclose( + target_accel_diff, target_accel[1:], atol=TOLERANCE) + np.testing.assert_allclose( + target_vel_diff, target_vel[1:], atol=TOLERANCE) # Make sure that the position and velocity targets are within bounds - robot = env.robot - pos_min = robot.position_limit_lower[robot.motor_position_indices[-1]] - pos_max = robot.position_limit_upper[robot.motor_position_indices[-1]] - vel_limit = robot.velocity_limit[robot.motor_velocity_indices[-1]] + motor = env.robot.motors[-1] self.assertTrue(np.all(np.logical_and( - pos_min <= target_pos, target_pos <= pos_max))) - self.assertTrue(np.all(np.abs(target_vel) <= vel_limit)) + motor.position_limit_lower <= target_pos, + target_pos <= motor.position_limit_upper))) + self.assertTrue(np.all(np.abs(target_vel) <= motor.velocity_limit)) def test_repeatability(self): # Instantiate the environment diff --git a/python/gym_jiminy/unit_py/test_pipeline_design.py b/python/gym_jiminy/unit_py/test_pipeline_design.py index 1bbca3557..0c39e76aa 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_design.py +++ b/python/gym_jiminy/unit_py/test_pipeline_design.py @@ -6,15 +6,19 @@ import unittest import numpy as np +import gymnasium as gym from jiminy_py.robot import _gcd -from gym_jiminy.common.utils import build_pipeline, load_pipeline +from jiminy_py.log import extract_trajectory_from_log +from gym_jiminy.common.utils import ( + save_trajectory_to_hdf5, build_pipeline, load_pipeline) from gym_jiminy.common.bases import InterfaceJiminyEnv from gym_jiminy.common.blocks import PDController TOLERANCE = 1.0e-6 + class PipelineDesign(unittest.TestCase): """ TODO: Write documentation """ @@ -43,9 +47,9 @@ def setUp(self): update_ratio=2, kp=self.pid_kp, kd=self.pid_kd, - target_position_margin=0.0, - target_velocity_limit=float("inf"), - target_acceleration_limit=float("inf") + joint_position_margin=0.0, + joint_velocity_limit=float("inf"), + joint_acceleration_limit=float("inf") ) ), wrapper=dict( @@ -78,7 +82,7 @@ def setUp(self): ) ), dict( wrapper=dict( - cls='gym_jiminy.common.wrappers.StackedJiminyEnv', + cls='gym_jiminy.common.wrappers.StackObservation', kwargs=dict( nested_filter_keys=[ ('t',), @@ -99,6 +103,16 @@ def test_load_files(self): # Get data path data_dir = os.path.join(os.path.dirname(__file__), "data") + # Generate machine-dependent reference trajectory + env = self.ANYmalPipelineEnv() + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + env.stop() + trajectory = extract_trajectory_from_log(env.log_data) + save_trajectory_to_hdf5( + trajectory, os.path.join(data_dir, "anymal_trajectory.hdf5")) + # Load TOML pipeline description, create env and perform a step toml_file = os.path.join(data_dir, "anymal_pipeline.toml") ANYmalPipelineEnv = load_pipeline(toml_file) diff --git a/python/gym_jiminy/unit_py/test_quantities.py b/python/gym_jiminy/unit_py/test_quantities.py index 9ae31c239..3a9bc8165 100644 --- a/python/gym_jiminy/unit_py/test_quantities.py +++ b/python/gym_jiminy/unit_py/test_quantities.py @@ -5,17 +5,23 @@ import numpy as np import gymnasium as gym + import jiminy_py +from jiminy_py.log import extract_trajectory_from_log import pinocchio as pin +from gym_jiminy.common.bases import QuantityEvalMode, DatasetTrajectoryQuantity from gym_jiminy.common.quantities import ( - QuantityManager, EulerAnglesFrame, CenterOfMass, ZeroMomentPoint) + QuantityManager, FrameEulerAngles, FrameXYZQuat, MaskedQuantity, + AverageFrameSpatialVelocity, CenterOfMass, ZeroMomentPoint) class Quantities(unittest.TestCase): """ TODO: Write documentation """ def test_shared_cache(self): + """ TODO: Write documentation + """ env = gym.make("gym_jiminy.envs:atlas") env.reset() @@ -25,7 +31,7 @@ def test_shared_cache(self): ("com", CenterOfMass, {}), ("zmp", ZeroMomentPoint, {})): quantity_manager[name] = (cls, kwargs) - quantities = quantity_manager._quantities + quantities = quantity_manager.registry assert len(quantity_manager) == 3 assert len(quantities["zmp"].cache.owners) == 1 @@ -53,21 +59,26 @@ def test_shared_cache(self): assert np.any(zmp_1 != zmp_2) def test_dynamic_batching(self): + """ TODO: Write documentation + """ env = gym.make("gym_jiminy.envs:atlas") env.reset() env.step(env.action) quantity_manager = QuantityManager(env) for name, cls, kwargs in ( - ("rpy_0", EulerAnglesFrame, dict( + ("xyzquat_0", FrameXYZQuat, dict( + frame_name=env.robot.pinocchio_model.frames[2].name)), + ("rpy_0", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[1].name)), - ("rpy_1", EulerAnglesFrame, dict( + ("rpy_1", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[1].name)), - ("rpy_2", EulerAnglesFrame, dict( + ("rpy_2", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[-1].name))): quantity_manager[name] = (cls, kwargs) - quantities = quantity_manager._quantities + quantities = quantity_manager.registry + xyzquat_0 = quantity_manager.xyzquat_0.copy() rpy_0 = quantity_manager.rpy_0.copy() assert len(quantities['rpy_0'].requirements['data'].frame_names) == 1 assert np.all(rpy_0 == quantity_manager.rpy_1) @@ -78,7 +89,9 @@ def test_dynamic_batching(self): env.step(env.action) quantity_manager.reset() rpy_0_next = quantity_manager.rpy_0 + xyzquat_0_next = quantity_manager.xyzquat_0.copy() assert np.any(rpy_0 != rpy_0_next) + assert np.any(xyzquat_0 != xyzquat_0_next) assert len(quantities['rpy_2'].requirements['data'].frame_names) == 2 assert len(quantities['rpy_1'].requirements['data'].cache.owners) == 3 @@ -90,22 +103,25 @@ def test_dynamic_batching(self): quantity_manager.reset(reset_tracking=True) assert np.all(rpy_0_next == quantity_manager.rpy_0) + assert np.all(xyzquat_0_next == quantity_manager.xyzquat_0) assert len(quantities['rpy_0'].requirements['data'].frame_names) == 1 def test_discard(self): + """ TODO: Write documentation + """ env = gym.make("gym_jiminy.envs:atlas") env.reset() quantity_manager = QuantityManager(env) for name, cls, kwargs in ( - ("rpy_0", EulerAnglesFrame, dict( + ("rpy_0", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[1].name)), - ("rpy_1", EulerAnglesFrame, dict( + ("rpy_1", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[1].name)), - ("rpy_2", EulerAnglesFrame, dict( + ("rpy_2", FrameEulerAngles, dict( frame_name=env.robot.pinocchio_model.frames[-1].name))): quantity_manager[name] = (cls, kwargs) - quantities = quantity_manager._quantities + quantities = quantity_manager.registry assert len(quantities['rpy_1'].cache.owners) == 2 assert len(quantities['rpy_2'].requirements['data'].cache.owners) == 3 @@ -121,17 +137,102 @@ def test_discard(self): del quantity_manager['rpy_2'] gc.collect() - for cache in quantity_manager._caches.values(): - assert len(cache.owners) == 0 + for (cls, _), cache in quantity_manager._caches.items(): + assert len(cache.owners) == (cls is DatasetTrajectoryQuantity) def test_env(self): + """ TODO: Write documentation + """ + env = gym.make("gym_jiminy.envs:atlas") + + env.quantities["zmp"] = (ZeroMomentPoint, {}) + + env.reset(seed=0) + zmp_0 = env.quantities["zmp"].copy() + env.step(env.action) + assert np.all(zmp_0 != env.quantities["zmp"]) + env.reset(seed=0) + assert np.all(zmp_0 == env.quantities["zmp"]) + + def test_stack(self): + """ TODO: Write documentation + """ env = gym.make("gym_jiminy.envs:atlas") + env.reset() - env.quantities["com"] = (CenterOfMass, {}) + quantity_cls = AverageFrameSpatialVelocity + quantity_kwargs = dict( + frame_name=env.robot.pinocchio_model.frames[1].name) + env.quantities["v_avg"] = (quantity_cls, quantity_kwargs) env.reset(seed=0) - com_0 = env.quantities["com"].copy() + with self.assertRaises(ValueError): + env.quantities["v_avg"] + + env.step(env.action) + v_avg = env.quantities["v_avg"].copy() + env.step(env.action) env.step(env.action) - assert np.all(com_0 != env.quantities["com"]) + assert np.all(v_avg != env.quantities["v_avg"]) + + def test_masked(self): + """ TODO: Write documentation + """ + env = gym.make("gym_jiminy.envs:atlas") + env.reset() + env.step(env.action) + + # 1. From non-slice-able indices + env.quantities["v_masked"] = (MaskedQuantity, dict( + quantity=(FrameXYZQuat, dict(frame_name="root_joint")), + key=(0, 1, 5))) + quantity = env.quantities.registry["v_masked"] + assert not quantity._slices + np.testing.assert_allclose( + env.quantities["v_masked"], quantity.data[[0, 1, 5]]) + del env.quantities["v_masked"] + + # 2. From boolean mask + env.quantities["v_masked"] = (MaskedQuantity, dict( + quantity=(FrameXYZQuat, dict(frame_name="root_joint")), + key=(True, True, False, False, False, True))) + quantity = env.quantities.registry["v_masked"] + np.testing.assert_allclose( + env.quantities["v_masked"], quantity.data[[0, 1, 5]]) + del env.quantities["v_masked"] + + # 3. From slice-able indices + env.quantities["v_masked"] = (MaskedQuantity, dict( + quantity=(FrameXYZQuat, dict(frame_name="root_joint")), + key=(0, 2, 4))) + quantity = env.quantities.registry["v_masked"] + assert len(quantity._slices) == 1 and quantity._slices[0] == slice(0, 5, 2) + np.testing.assert_allclose( + env.quantities["v_masked"], quantity.data[[0, 2, 4]]) + + def test_true_vs_reference(self): + env = gym.make("gym_jiminy.envs:atlas") + + env.quantities["zmp"] = ( + ZeroMomentPoint, dict(mode=QuantityEvalMode.TRUE)) + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + zmp_0 = env.quantities["zmp"].copy() + env.stop() + + trajectory = extract_trajectory_from_log(env.log_data) + env.quantities["zmp_ref"] = ( + ZeroMomentPoint, dict(mode=QuantityEvalMode.REFERENCE)) + + with self.assertRaises(RuntimeError): + env.reset(seed=0) + + env.quantities.add_trajectory("reference", trajectory) + env.quantities.select_trajectory("reference") + env.reset(seed=0) - assert np.all(com_0 == env.quantities["com"]) + for _ in range(10): + env.step(env.action_space.sample() * 0.05) + assert np.all(zmp_0 != env.quantities["zmp"]) + np.testing.assert_allclose(zmp_0, env.quantities["zmp_ref"]) diff --git a/python/gym_jiminy/unit_py/test_reward.py b/python/gym_jiminy/unit_py/test_reward.py new file mode 100644 index 000000000..2dbed8f19 --- /dev/null +++ b/python/gym_jiminy/unit_py/test_reward.py @@ -0,0 +1,95 @@ +""" TODO: Write documentation +""" +import unittest + +import numpy as np + +import gymnasium as gym +import jiminy_py +import pinocchio as pin +from jiminy_py.log import extract_trajectory_from_log + +from gym_jiminy.common.compositions import ( + OdometryVelocityReward, SurviveReward, AdditiveMixtureReward) + + +class Rewards(unittest.TestCase): + """ TODO: Write documentation + """ + def test_average_odometry_velocity(self): + """ TODO: Write documentation + """ + env = gym.make("gym_jiminy.envs:atlas") + + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + env.stop() + trajectory = extract_trajectory_from_log(env.log_data) + env.quantities.add_trajectory("reference", trajectory) + env.quantities.select_trajectory("reference") + + reward = OdometryVelocityReward(env, cutoff=0.3) + quantity_odom_vel_true = reward.quantity.requirements['value_left'] + quantity_odom_vel_ref = reward.quantity.requirements['value_left'] + + env.reset(seed=0) + base_pose_prev = env.robot_state.q[:7].copy() + _, _, terminated, _, _ = env.step(env.action) + base_pose = env.robot_state.q[:7].copy() + + se3 = pin.liegroups.SE3() + base_pose_diff = pin.LieGroup.difference( + se3, base_pose_prev, base_pose) + base_velocity_mean_local = base_pose_diff / env.step_dt + base_pose_mean = pin.LieGroup.integrate( + se3, base_pose_prev, 0.5 * base_pose_diff) + rot_mat = pin.Quaternion(base_pose_mean[-4:]).matrix() + base_velocity_mean_world = np.concatenate(( + rot_mat @ base_velocity_mean_local[:3], + rot_mat @ base_velocity_mean_local[3:])) + np.testing.assert_allclose( + quantity_odom_vel_true.requirements['data'].data, + base_velocity_mean_world) + base_odom_velocity = base_velocity_mean_world[[0, 1, 5]] + np.testing.assert_allclose( + quantity_odom_vel_true.data, base_odom_velocity) + gamma = - np.log(0.01) / reward.cutoff ** 2 + value = np.exp(- gamma * np.sum(( + quantity_odom_vel_true.data - quantity_odom_vel_ref.data) ** 2)) + np.testing.assert_allclose(reward(terminated, {}), value) + + def test_mixture(self): + env = gym.make("gym_jiminy.envs:atlas") + + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + env.stop() + trajectory = extract_trajectory_from_log(env.log_data) + env.quantities.add_trajectory("reference", trajectory) + env.quantities.select_trajectory("reference") + + reward_odometry = OdometryVelocityReward(env, cutoff=0.3) + reward_survive = SurviveReward(env) + reward_sum = AdditiveMixtureReward( + env, + "reward_total", + components=(reward_odometry, reward_survive), + weights=(0.5, 0.2)) + reward_sum_normalized = AdditiveMixtureReward( + env, + "reward_total", + components=(reward_odometry, reward_survive), + weights=(0.7, 0.3)) + + env.reset(seed=0) + _, _, terminated, _, _ = env.step(env.action) + + assert len(reward_sum_normalized.components) == 2 + assert reward_sum_normalized.is_terminal == False + assert reward_sum_normalized.is_normalized + assert not reward_sum.is_normalized + assert reward_sum(terminated, {}) == ( + 0.5 * reward_odometry(terminated, {}) + + 0.2 * reward_survive(terminated, {})) diff --git a/python/jiminy_py/examples/double_pendulum.py b/python/jiminy_py/examples/double_pendulum.py index 788fd03ac..0a1cd291f 100644 --- a/python/jiminy_py/examples/double_pendulum.py +++ b/python/jiminy_py/examples/double_pendulum.py @@ -52,9 +52,9 @@ def internal_dynamics(self, t, q, v, u_custom): engine_options["telemetry"]["enableConfiguration"] = True engine_options["telemetry"]["enableVelocity"] = True engine_options["telemetry"]["enableAcceleration"] = True + engine_options["telemetry"]["enableEffort"] = True engine_options["telemetry"]["enableForceExternal"] = False engine_options["telemetry"]["enableCommand"] = True - engine_options["telemetry"]["enableMotorEffort"] = True engine_options["telemetry"]["enableEnergy"] = True engine_options["world"]["gravity"][2] = -9.81 engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" # ["runge_kutta_dopri5", "euler_explicit"] @@ -107,8 +107,9 @@ def internal_dynamics(self, t, q, v, u_custom): plt.grid() plt.show() - # Display the simulation trajectory and the reference - play_logs_data(robot, log_data, + # Replay the simulation + play_logs_data(log_data, + robot, speed_ratio=0.5, camera_pose=((0.0, 7.0, 0.0), (np.pi/2, 0.0, np.pi), None), delete_robot_on_close=False) diff --git a/python/jiminy_py/setup.py b/python/jiminy_py/setup.py index d432b24cc..2746ecbcd 100644 --- a/python/jiminy_py/setup.py +++ b/python/jiminy_py/setup.py @@ -144,7 +144,7 @@ def finalize_options(self) -> None: # Used internally by Viewer to enable recording video # programmatically with Meshcat as backend. # - 1.43 is broken - "playwright!=1.43" + "playwright<1.43" ], "dev": [ # Generate Python type hints files (aka. stubs) for C extensions. @@ -156,7 +156,7 @@ def finalize_options(self) -> None: # Used in unit tests for numerical integration and interpolation "scipy", # Stub for static type checking - "types-psutil", + "types-psutil>=5.9.5.20240511", "types-Pillow", "types-toml", "types-tqdm", diff --git a/python/jiminy_py/src/jiminy_py/core/__init__.py b/python/jiminy_py/src/jiminy_py/core/__init__.py index c5c922e6f..0b3e1a415 100644 --- a/python/jiminy_py/src/jiminy_py/core/__init__.py +++ b/python/jiminy_py/src/jiminy_py/core/__init__.py @@ -107,19 +107,37 @@ # Define helpers to build extension modules def get_cmake_module_path() -> str: - """ TODO: Write documentation. + """Get Jiminy Core cmake module path. + + One must add this path to CMake search path for `find_package` to be able + to detect Jiminy core install. + + include it when building C++ extension modules if Jiminy has + been installed via `pip`. """ return _os.path.join(_os.path.dirname(__file__), "cmake") def get_include() -> str: - """ TODO: Write documentation. + """Get Jiminy Core include directory. + + One must include it when building C++ extension modules if Jiminy has been + installed via `pip`. This can be done automatically when CMake build system + is used. All it requires is importing Jiminy C++ API via `find_package`, + then declaring `jiminy::core` imported target as target link library of the + extension module using `target_link_libraries`. """ return _os.path.join(_os.path.dirname(__file__), "include") def get_libraries() -> str: - """ TODO: Write documentation. + """List all Jiminy Core shared libraries. + + One must include it when building C++ extension modules if Jiminy has been + installed via `pip`. This can be done automatically when CMake build system + is used. All it requires is importing Jiminy C++ API via `find_package`, + then declaring `jiminy::core` imported target as target link library of the + extension module using `target_link_libraries`. """ ver_short = '.'.join(__version__.split('.')[:2]) lib_dir = _os.path.join(_os.path.dirname(__file__), "lib") diff --git a/python/jiminy_py/src/jiminy_py/dynamics.py b/python/jiminy_py/src/jiminy_py/dynamics.py index 628636b89..373655a6a 100644 --- a/python/jiminy_py/src/jiminy_py/dynamics.py +++ b/python/jiminy_py/src/jiminy_py/dynamics.py @@ -1,10 +1,15 @@ # mypy: disable-error-code="attr-defined, name-defined" """Helpers to ease computation of kinematic and dynamic quantities. + +.. warning:: + These helpers must be used with caution. They are inefficient and some may + not even work properly due to ground profile being partially supported. """ # pylint: disable=invalid-name,no-member import logging -from copy import deepcopy -from typing import Optional, Tuple, Sequence, Callable, TypedDict, Any +from bisect import bisect_right +from dataclasses import dataclass +from typing import Optional, Tuple, Sequence, Callable, Literal import numpy as np @@ -22,6 +27,9 @@ LOGGER = logging.getLogger(__name__) +TRAJ_INTERP_TOL = 1e-12 # 0.01 * 'STEPPER_MIN_TIMESTEP' + + # ##################################################################### # ######################### Generic math ############################## # ##################################################################### @@ -88,69 +96,238 @@ def velocityXYZQuatToXYZRPY(xyzquat: np.ndarray, # #################### State and Trajectory ########################### # ##################################################################### +@dataclass(unsafe_hash=True) class State: - """Store the kinematics and dynamics data of the robot at a given time. + """Basic data structure storing kinematics and dynamics information at a + given time. + + .. note:: + The user is the responsible for keeping track to which robot the state + is associated to as this information is not stored in the state itself. + """ + + t: float + """Time. + """ + + q: np.ndarray + """Configuration vector. + """ + + v: Optional[np.ndarray] = None + """Velocity vector as a 1D array. + """ + + a: Optional[np.ndarray] = None + """Acceleration vector as a 1D array. + """ + + u: Optional[np.ndarray] = None + """Total joint efforts as a 1D array. + """ + + command: Optional[np.ndarray] = None + """Motor command as a 1D array. + """ + + f_external: Optional[np.ndarray] = None + """Joint external forces as a 2D array. + + The first dimension corresponds to the N individual joints of the robot + including 'universe', while the second gathers the 6 spatial force + coordinates (Fx, Fy, Fz, Mx, My, Mz). + """ + + +@dataclass(unsafe_hash=True) +class Trajectory: + """Trajectory of a robot. + + This class is mostly a basic data structure storing a sequence of states + along with the robot to which it is associated. On top of that, it provides + helper methods to make it easier to manipulate these data, eg query the + state at a given timestamp. """ - def __init__(self, # pylint: disable=unused-argument - t: float, - q: np.ndarray, - v: Optional[np.ndarray] = None, - a: Optional[np.ndarray] = None, - tau: Optional[np.ndarray] = None, - contact_frames: Optional[Sequence[str]] = None, - f_ext: Optional[Sequence[np.ndarray]] = None, - copy: bool = False, - **kwargs: Any) -> None: + + states: Tuple[State, ...] + """Sequence of states of increasing time. + + .. warning:: + The time may not be strictly increasing. There may be up to two + consecutive data point associated with the same timestep because + quantities may vary instantaneously at acceleration-level and higher. + """ + + robot: jiminy.Robot + """Robot associated with the trajectory. + """ + + use_theoretical_model: bool + """Whether to use the theoretical model or the extended simulation model. + """ + + def __init__(self, + states: Sequence[State], + robot: jiminy.Robot, + use_theoretical_model: bool) -> None: """ - :param t: Time. - :param q: Configuration vector. - :param v: Velocity vector. - :param a: Acceleration vector. - :param tau: Joint efforts. - :param contact_frames: Name of the contact frames. - :param f_ext: Joint external forces. - :param copy: Force to copy the arguments. + :param states: Trajectory data as a sequence of `State` instances of + increasing time. + :param robot: Robot associated with the trajectory. + :param use_theoretical_model: Whether the state sequence is associated + with the theoretical dynamical model or + extended simulation model of the robot. """ - # Time - self.t = t - # Configuration vector - self.q = deepcopy(q) if copy else q - # Velocity vector - self.v = deepcopy(v) if copy else v - # Acceleration vector - self.a = deepcopy(a) if copy else a - # Effort vector - self.tau = deepcopy(tau) if copy else tau - # Frame names of the contact points - if copy: - self.contact_frames = deepcopy(contact_frames) + # Backup user arguments + self.states = tuple(states) + self.robot = robot + self.use_theoretical_model = use_theoretical_model + + # Extract time associated with all states + self._times = tuple(state.t for state in states) + if any(t_right - t_left < 0.0 for t_right, t_left in zip( + self._times[1:], self._times[:-1])): + raise ValueError( + "Time must not be decreasing between consecutive timesteps.") + + # Define pinocchio model proxy for fast access + if use_theoretical_model: + self._pinocchio_model = robot.pinocchio_model_th else: - self.contact_frames = contact_frames - # External forces - self.f_ext = deepcopy(f_ext) if copy else f_ext + self._pinocchio_model = robot.pinocchio_model + + # Keep track of last request to speed up nearest neighbors search + self._t_prev = 0.0 + self._index_prev = 0 + + # List of optional state fields that are provided + self._fields: Tuple[str, ...] = () + self._has_velocity = False + self._has_acceleration = False + self._has_effort = False + self._has_command = False + self._has_external_forces = False + if states: + state = states[0] + self._has_velocity = state.v is not None + self._has_acceleration = state.a is not None + self._has_effort = state.u is not None + self._has_command = state.command is not None + self._has_external_forces = state.f_external is not None + self._fields = tuple( + field for field in ("v", "a", "u", "command", "f_external") + if getattr(state, field) is not None) + + @property + def has_data(self) -> bool: + """Whether the trajectory has data, ie the state sequence is not empty. + """ + return bool(self.states) - def __repr__(self) -> str: - """Convert the kinematics and dynamics data into string. + @property + def has_velocity(self) -> bool: + """Whether the trajectory contains the velocity vector. + """ + return self._has_velocity - :returns: The kinematics and dynamics data as a string. + @property + def has_acceleration(self) -> bool: + """Whether the trajectory contains the acceleration vector. """ - msg = "" - for key, val in self.__dict__.items(): - if val is not None: - msg += f"{key} : {val}\n" - return msg + return self._has_acceleration + @property + def has_effort(self) -> bool: + """Whether the trajectory contains the effort vector. + """ + return self._has_acceleration -class TrajectoryDataType(TypedDict, total=False): - """Basic data structure storing the required information about a trajectory - to later replay it using `jiminy_py.viewer.play_trajectories`. - """ - # List of State objects of increasing time - evolution_robot: Sequence[State] - # Jiminy robot - robot: jiminy.Robot - # Whether to use the theoretical model or the extended simulation model - use_theoretical_model: bool + @property + def has_command(self) -> bool: + """Whether the trajectory contains motor commands. + """ + return self._has_command + + @property + def has_external_forces(self) -> bool: + """Whether the trajectory contains external forces. + """ + return self._has_external_forces + + @property + def time_interval(self) -> Tuple[float, float]: + """Time interval of the trajectory. + + It raises an exception if no data is available. + """ + if not self.has_data: + raise RuntimeError( + "State sequence is empty. Time interval undefined.") + return (self._times[0], self._times[-1]) + + def get(self, + t: float, + mode: Literal['raise', 'wrap', 'clip'] = 'raise') -> State: + """Query the state at a given timestamp. + + Internally, the nearest neighbor states are linearly interpolated, + taking into account the corresponding Lie Group of all state attributes + that are available. + + :param t: Time of the state to extract from the trajectory. + :param mode: Specifies how to deal with query time of are out of the + time interval 'time_interval' of the trajectory. Specify + 'raise' to raise an exception if the query time is + out-of-bound wrt to underlying state sequence of the + selected trajectory. Specify 'clip' to force clipping of + the query time before interpolation of the state sequence. + Specify 'wrap' to wrap around the query time wrt the time + span of the trajectory. This is useful to store periodic + trajectories as finite state sequences. + """ + # Raise exception if state sequence is empty + if not self.has_data: + raise RuntimeError( + "State sequence is empty. Impossible to interpolate data.") + + # Backup the original query time + t_orig = t + + # Handling of the desired mode + t_start, t_end = self.time_interval + if mode == "raise": + if t - t_end > TRAJ_INTERP_TOL or t_start - t > TRAJ_INTERP_TOL: + raise RuntimeError("Time is out-of-range.") + elif mode == "wrap": + t = ((t - t_start) % (t_end - t_start)) + t_start + else: + t = max(t, t_start) # Clipping right it is sufficient + + # Get nearest neighbors timesteps for linear interpolation + if t < self._t_prev: + self._index_prev = 0 + self._index_prev = bisect_right( + self._times, t, self._index_prev, len(self._times) - 1) + self._t_prev = t + + # Skip interpolation if not necessary + index_left, index_right = self._index_prev - 1, self._index_prev + t_left, s_left = self._times[index_left], self.states[index_left] + if t - t_left < TRAJ_INTERP_TOL: + return s_left + t_right, s_right = self._times[index_right], self.states[index_right] + if t_right - t < TRAJ_INTERP_TOL: + return s_right + alpha = (t - t_left) / (t_right - t_left) + + # Interpolate state + data = {"q": pin.interpolate( + self._pinocchio_model, s_left.q, s_right.q, alpha)} + for field in self._fields: + value_left = getattr(s_left, field) + value_right = getattr(s_right, field) + data[field] = value_left + alpha * (value_right - value_left) + return State(t=t_orig, **data) # ##################################################################### @@ -162,7 +339,7 @@ def update_quantities(robot: jiminy.Model, velocity: Optional[np.ndarray] = None, acceleration: Optional[np.ndarray] = None, update_physics: bool = True, - update_com: bool = True, + update_centroidal: bool = True, update_energy: bool = True, update_jacobian: bool = False, update_collisions: bool = True, @@ -210,9 +387,9 @@ def update_quantities(robot: jiminy.Model, :param update_physics: Whether to compute the non-linear effects and internal/external forces. Optional: True by default. - :param update_com: Whether to compute the COM of the robot AND each link - individually. The global COM is the first index. - Optional: False by default. + :param update_centroidal: Whether to compute the centroidal dynamics (incl. + CoM) of the robot. + Optional: False by default. :param update_energy: Whether to compute the energy of the robot. Optional: False by default :param update_jacobian: Whether to compute the jacobians. @@ -229,7 +406,7 @@ def update_quantities(robot: jiminy.Model, model = robot.pinocchio_model data = robot.pinocchio_data - if (update_physics and update_com and + if (update_physics and update_centroidal and update_energy and update_jacobian and velocity is not None and acceleration is None): pin.computeAllTerms(model, data, position, velocity) @@ -242,7 +419,7 @@ def update_quantities(robot: jiminy.Model, pin.forwardKinematics( model, data, position, velocity, acceleration) - if update_com: + if update_centroidal: if velocity is None: kinematic_level = pin.POSITION elif acceleration is None: @@ -250,9 +427,10 @@ def update_quantities(robot: jiminy.Model, else: kinematic_level = pin.ACCELERATION pin.centerOfMass(model, data, kinematic_level, False) + pin.computeCentroidalMomentumTimeVariation(model, data) if update_jacobian: - if update_com: + if update_centroidal: pin.jacobianCenterOfMass(model, data) pin.computeJointJacobians(model, data) @@ -565,7 +743,7 @@ def compute_freeflyer_state_from_fixed_body( :param acceleration: See position. :param fixed_body_name: Name of the body frame that is considered fixed parallel to world frame. - Optional: It will be infered from the set of + Optional: It will be inferred from the set of contact points and collision bodies. :param ground_profile: Ground profile callback. :param use_theoretical_model: @@ -599,7 +777,7 @@ def compute_freeflyer_state_from_fixed_body( velocity, acceleration, update_physics=False, - update_com=False, + update_centroidal=False, update_energy=False, use_theoretical_model=use_theoretical_model) @@ -684,10 +862,10 @@ def compute_efforts_from_fixed_body( data.oMi[1]).act(data.f[1]) # Recompute the efforts with RNEA and the correct external forces - tau = jiminy.rnea(model, data, position, velocity, acceleration, f_ext) + u = jiminy.rnea(model, data, position, velocity, acceleration, f_ext) f_ext = f_ext[support_joint_index] - return tau, f_ext + return u, f_ext def compute_inverse_dynamics(robot: jiminy.Model, @@ -697,7 +875,7 @@ def compute_inverse_dynamics(robot: jiminy.Model, use_theoretical_model: bool = False ) -> np.ndarray: """Compute the motor torques through inverse dynamics, assuming to external - forces except the one resulting from the anyaltical constraints applied on + forces except the one resulting from the analytical constraints applied on the model. .. warning:: @@ -728,7 +906,8 @@ def compute_inverse_dynamics(robot: jiminy.Model, # Define some proxies for convenience model = robot.pinocchio_model data = robot.pinocchio_data - motor_velocity_indices = robot.motor_velocity_indices + motor_velocity_indices = [ + model.joints[motor.joint_index].idx_v for motor in robot.motors] # Updating kinematics quantities pin.forwardKinematics(model, data, position, velocity, acceleration) @@ -761,63 +940,3 @@ def compute_inverse_dynamics(robot: jiminy.Model, u = eigenpy.LDLT(B_ydd).solve(- a_ydd) return u - - -# ##################################################################### -# ################### State sequence wrappers ######################### -# ##################################################################### - -def compute_freeflyer(trajectory_data: TrajectoryDataType, - freeflyer_continuity: bool = True) -> None: - """Compute the freeflyer positions and velocities. - - .. warning:: - This function modifies the internal robot data. - - :param trajectory_data: Sequence of States for which to retrieve the - freeflyer. - :param freeflyer_continuity: Whether to enforce the continuity in position - of the freeflyer. - Optional: True by default. - """ - robot = trajectory_data['robot'] - - contact_frame_prev: Optional[str] = None - w_M_ff_offset = pin.SE3.Identity() - w_M_ff_prev = None - for s in trajectory_data['evolution_robot']: - # Compute freeflyer using contact frame as reference frame - compute_freeflyer_state_from_fixed_body( - robot, s.q, s.v, s.a, s.contact_frame, None) - - # Move freeflyer to ensure continuity over time, if requested - if freeflyer_continuity: - # Extract the current freeflyer transform - w_M_ff = pin.XYZQUATToSE3(s.q[:7]) - - # Update the internal buffer of the freeflyer transform - if (contact_frame_prev is not None and - contact_frame_prev != s.contact_frame): - w_M_ff_offset = w_M_ff_offset * w_M_ff_prev * w_M_ff.inverse() - contact_frame_prev = s.contact_frame - w_M_ff_prev = w_M_ff - - # Add the appropriate offset to the freeflyer - w_M_ff = w_M_ff_offset * w_M_ff - s.q[:7] = pin.SE3ToXYZQUAT(w_M_ff) - - -def compute_efforts(trajectory_data: TrajectoryDataType) -> None: - """Compute the efforts in the trajectory using RNEA method. - - :param trajectory_data: Sequence of States for which to compute the - efforts. - """ - robot = trajectory_data['robot'] - use_theoretical_model = trajectory_data['use_theoretical_model'] - - for s in trajectory_data['evolution_robot']: - assert s.q is not None and s.v is not None and s.a is not None - assert s.contact_frames is not None - s.tau, s.f_ext = compute_efforts_from_fixed_body( - robot, s.q, s.v, s.a, s.contact_frame, use_theoretical_model) diff --git a/python/jiminy_py/src/jiminy_py/log.py b/python/jiminy_py/src/jiminy_py/log.py index 2256714df..fc8f2a472 100644 --- a/python/jiminy_py/src/jiminy_py/log.py +++ b/python/jiminy_py/src/jiminy_py/log.py @@ -4,39 +4,17 @@ """ import re from bisect import bisect_right +from itertools import zip_longest, starmap from collections import OrderedDict from typing import ( - Any, Callable, List, Dict, Optional, Sequence, Union, Literal, Type, - overload) + Any, Callable, List, Dict, Optional, Sequence, Union, Literal, overload, + cast) import numpy as np from . import core as jiminy from . import tree -from .core import ( # pylint: disable=no-name-in-module - EncoderSensor as encoder, - EffortSensor as effort, - ContactSensor as contact, - ForceSensor as force, - ImuSensor as imu) -from .dynamics import State, TrajectoryDataType - - -SENSORS_FIELDS: Dict[ - Type[jiminy.AbstractSensor], Union[List[str], Dict[str, List[str]]] - ] = { - encoder: encoder.fieldnames, - effort: effort.fieldnames, - contact: contact.fieldnames, - force: { - k: [e[len(k):] for e in force.fieldnames if e.startswith(k)] - for k in ['F', 'M'] - }, - imu: { - k: [e[len(k):] for e in imu.fieldnames if e.startswith(k)] - for k in ['Quat', 'Gyro', 'Accel'] - } -} +from .dynamics import State, Trajectory FieldNested = Union[Dict[str, 'FieldNested'], Sequence['FieldNested'], str] @@ -96,7 +74,7 @@ def extract_variables_from_log(log_vars: Dict[str, np.ndarray], # Raise an exception if the key does not exists and not fail safe if key not in log_vars: - raise ValueError(f"Variable '{key}' not found in log file.") + raise KeyError(f"Variable '{key}' not found in log file.") # Extract the value corresponding to the key if as_dict: @@ -168,8 +146,8 @@ def build_robot_from_log( # Get binary serialized robot data robot_data = log_constants[".".join(filter(None, (robot_name, "robot")))] - # Load robot - return jiminy.build_robot_from_binary( + # Load robot from binary data + return jiminy.load_from_binary( robot_data, mesh_path_dir, mesh_package_dirs) @@ -215,7 +193,7 @@ def build_robots_from_log( def extract_trajectory_from_log(log_data: Dict[str, Any], robot: Optional[jiminy.Robot] = None, *, robot_name: Optional[str] = None - ) -> TrajectoryDataType: + ) -> Trajectory: """Extract the minimal required information from raw log data in order to replay the simulation in a viewer. @@ -232,9 +210,8 @@ def extract_trajectory_from_log(log_data: Dict[str, Any], :param robot_name: Name of the robot to be constructed in the log. If it is not known, call `build_robot_from_log`. - :returns: Trajectory dictionary. The actual trajectory corresponds to the - field "evolution_robot" and it is a list of State object. The - other fields are additional information. + :returns: Dictionary whose keys are the name of each robot and values are + the corresponding `Trajectory` object. """ # Prepare robot namespace if robot is not None: @@ -254,45 +231,44 @@ def extract_trajectory_from_log(log_data: Dict[str, Any], # Define some proxies for convenience log_vars = log_data["variables"] - # Extract the joint positions, velocities and external forces over time - positions = np.stack([ - log_vars.get(".".join(filter(None, (robot_name, field))), []) - for field in robot.log_position_fieldnames], axis=-1) - velocities = np.stack([ - log_vars.get(".".join(filter(None, (robot_name, field))), []) - for field in robot.log_velocity_fieldnames], axis=-1) - forces = np.stack([ - log_vars.get(".".join(filter(None, (robot_name, field))), []) - for field in robot.log_f_external_fieldnames], axis=-1) - - # Determine which optional data are available - has_positions = len(positions) > 0 - has_velocities = len(velocities) > 0 - has_forces = len(forces) > 0 + # Extract robot state data over time for all quantities available + data: Dict[str, Union[Sequence[np.ndarray], np.ndarray]] = OrderedDict() + for name in ("position", + "velocity", + "acceleration", + "effort", + "command", + "f_external"): + fieldnames = getattr(robot, f"log_{name}_fieldnames") + try: + data[name] = extract_variables_from_log( + log_vars, fieldnames, robot_name) + except KeyError: + data[name] = [] + + # Add fictitious 'universe' external force and reshape data if available + f_ext: np.ndarray = cast(np.ndarray, data["f_external"]) + if len(f_ext) > 0: + f_ext = np.stack((*((np.zeros_like(f_ext[0]),) * 6), *f_ext), axis=-1) + data["f_external"] = f_ext.reshape((len(f_ext), -1, 6), order='A') + + # Stack available data + for key, values in data.items(): + if len(values) > 0 and not isinstance(values, np.ndarray): + data[key] = np.stack(values, axis=-1) # Create state sequence - evolution_robot = [] - q, v, f_ext = None, None, None - for i, t in enumerate(log_vars["Global.Time"]): - if has_positions: - q = positions[i] - if has_velocities: - v = velocities[i] - if has_forces: - f_ext = [forces[i, (6 * (j - 1)):(6 * j)] - for j in range(1, robot.pinocchio_model.njoints)] - evolution_robot.append(State( - t=t, q=q, v=v, f_ext=f_ext)) # type: ignore[arg-type] - - return {"evolution_robot": evolution_robot, - "robot": robot, - "use_theoretical_model": False} + states = tuple(starmap( + State, zip_longest(log_vars["Global.Time"], *data.values()))) + + # Create the trajectory + return Trajectory(states, robot, use_theoretical_model=False) def extract_trajectories_from_log( log_data: Dict[str, Any], robots: Optional[Sequence[jiminy.Robot]] = None - ) -> Dict[str, TrajectoryDataType]: + ) -> Dict[str, Trajectory]: """Extract the minimal required information from raw log data in order to replay the simulation in a viewer. @@ -302,12 +278,12 @@ def extract_trajectories_from_log( information. :param log_data: Logged data (constants and variables) as a dictionary. - :param robots: Sequend of Jiminy robots associated with the logged + :param robots: Sequence of Jiminy robots associated with the logged trajectory. Optional: None by default. If None, then it will be reconstructed from 'log_data' using `build_robots_from_log`. - :returns: Dictonary mapping each robot name to its corresponding + :returns: Dictionary mapping each robot name to its corresponding trajectory. """ # Handling of default argument(s) @@ -319,7 +295,6 @@ def extract_trajectories_from_log( for robot in robots: trajectories[robot.name] = extract_trajectory_from_log( log_data, robot, robot_name=robot.name) - return trajectories @@ -345,40 +320,41 @@ def update_sensor_measurements_from_log( times = log_vars["Global.Time"] # Filter sensors whose data is available - sensors_set, sensors_log = [], [] - for sensor_type, sensor_names in robot.sensor_names.items(): - sensor_fieldnames = getattr(jiminy, sensor_type).fieldnames - for name in sensor_names: - sensor = robot.get_sensor(sensor_type, name) - log_fieldnames = [ - '.'.join((sensor.name, field)) for field in sensor_fieldnames] - if log_fieldnames[0] in log_vars.keys(): - sensor_log = np.stack([ - log_vars[field] for field in log_fieldnames], axis=-1) - sensors_set.append(sensor) - sensors_log.append(sensor_log) + sensor_data_map = [] + for sensor_group in robot.sensors.values(): + for sensor in sensor_group: + log_sensor_name = ".".join((sensor.type, sensor.name)) + try: + data = np.stack(extract_variables_from_log( + log_vars, sensor.fieldnames, log_sensor_name), axis=-1) + sensor_data_map.append((sensor, data)) + except KeyError: + pass def update_hook(t: float, q: np.ndarray, # pylint: disable=unused-argument v: np.ndarray # pylint: disable=unused-argument ) -> None: - nonlocal times, sensors_set, sensors_log + nonlocal times, sensor_data_map + + # Early return if no more data is available + if t > times[-1]: + return # Get surrounding indices in log data i = bisect_right(times, t) i_prev, i_next = max(i - 1, 0), min(i, len(times) - 1) - # Early return if no more data is available + # Compute linear interpolation ratio if i_next == i_prev: - return - - # Compute current time ratio - ratio = (t - times[i_prev]) / (times[i_next] - times[i_prev]) + # Special case for final data point + ratio = 1.0 + else: + ratio = (t - times[i_prev]) / (times[i_next] - times[i_prev]) # Update sensors data - for sensor, sensor_log, in zip(sensors_set, sensors_log): - value_prev, value_next = sensor_log[i_prev], sensor_log[i_next] - value = value_prev + (value_next - value_prev) * ratio - sensor.data = value + for sensor, data in sensor_data_map: + value_prev, value_next = data[i_prev], data[i_next] + sensor.data = value_prev + (value_next - value_prev) * ratio return update_hook diff --git a/python/jiminy_py/src/jiminy_py/plot.py b/python/jiminy_py/src/jiminy_py/plot.py index 654cb562d..06a94f789 100644 --- a/python/jiminy_py/src/jiminy_py/plot.py +++ b/python/jiminy_py/src/jiminy_py/plot.py @@ -10,6 +10,7 @@ import pathlib import argparse import logging +from dataclasses import dataclass from math import ceil, sqrt, floor from textwrap import dedent from itertools import cycle @@ -17,7 +18,7 @@ from collections import OrderedDict from weakref import WeakKeyDictionary from typing import ( - Dict, Optional, Any, Tuple, List, Union, Callable, TypedDict, cast) + Dict, Any, List, Optional, Tuple, Union, Callable, Type, cast) import numpy as np try: @@ -41,13 +42,31 @@ from matplotlib.backends.backend_pdf import PdfPages import jiminy_py.core as jiminy -from .log import (SENSORS_FIELDS, - read_log, +from .core import ( # pylint: disable=no-name-in-module + EncoderSensor, EffortSensor, ContactSensor, ForceSensor, ImuSensor) +from .log import (read_log, extract_variables_from_log, build_robot_from_log) from .viewer import interactive_mode +SENSORS_FIELDS: Dict[ + Type[jiminy.AbstractSensor], Union[List[str], Dict[str, List[str]]] + ] = { + EncoderSensor: EncoderSensor.fieldnames, + EffortSensor: EffortSensor.fieldnames, + ContactSensor: ContactSensor.fieldnames, + ForceSensor: { + k: [e[len(k):] for e in ForceSensor.fieldnames if e.startswith(k)] + for k in ['F', 'M'] + }, + ImuSensor: { + k: [e[len(k):] for e in ImuSensor.fieldnames if e.startswith(k)] + for k in ['Quat', 'Gyro', 'Accel'] + } +} + + class _ButtonBlit(Button): def _motion(self, event: Event) -> None: if self.ignore(event): @@ -59,34 +78,57 @@ def _motion(self, event: Event) -> None: # It is necessary to flush events beforehand to make sure # figure refresh cannot get interrupted by button blitting. # Otherwise the figure would be blank. + # FIXME: `flush_events` on 'matplotlib>=3.8' causes deadlock assert self.ax.figure is not None - self.ax.figure.canvas.flush_events() + # self.ax.figure.canvas.flush_events() self.ax.draw_artist(self.ax) self.ax.figure.canvas.blit(self.ax.bbox) -class TabData(TypedDict, total=True): +@dataclass +class TabData: """Internal data stored for each tab if `TabbedFigure`. """ - # Matpolib `Axes` handles of every subplots + axes: List[Axes] - # Matpolib Legend data as returned by `get_legend_handles_labels` method. - # First element of pair is a list of Artist to legend, which usually are - # `line2D` object, and the second element is the list of labels for each - # of them. + """Matpolib `Axes` handles of every subplots. + """ + legend_data: Tuple[List[Artist], List[str]] - # Matplotlib `NavigationToolbar2` navigation stack history. It is not meant - # to be modified manually, but only copied/restored when needed. - nav_stack: List[WeakKeyDictionary] - # Matplotlib `NavigationToolbar2` navigation stack position, used - # internally for undo/redo history. It is not meant to be modified - # manually, but only copied/restored when needed. - nav_pos: int - # Button associated with the tab, on which to click to switch between tabs. + """Matpolib Legend data as returned by `get_legend_handles_labels` method. + + First element of pair is a list of Artist to legend, which usually are + `line2D` object, and the second element is the list of labels for each of + them. + """ + button: _ButtonBlit - # Axe of the button, used internally to define the position and size of - # the button. + """Button on which to click to switch to the tab at hands. + """ + button_axcut: Axes + """Axe of the button. + + .. note:: + This attribute is used internally to define the position and size of + the button. + """ + + nav_stack: List[WeakKeyDictionary] + """Matplotlib `NavigationToolbar2` navigation stack history. + + .. note:: + This attribute is used internally for undo/redo history. It is not + meant to be modified manually, but only copied/restored when needed. + """ + + nav_pos: int + """Matplotlib `NavigationToolbar2` navigation stack position. + + .. note:: + This attribute is used internally for undo/redo history. It is not + meant to be modified manually, but only copied/restored when needed. + """ class TabbedFigure: @@ -195,11 +237,11 @@ def adjust_layout(self, # Refresh button size, in case the number of tabs has changed buttons_width = (1.0 - 0.006) / len(self.tabs_data) for i, tab in enumerate(self.tabs_data.values()): - tab["button_axcut"].set_position( + tab.button_axcut.set_position( (buttons_width * i + 0.003, 0.1, buttons_width, 1.0)) # Re-arrange subplots in case figure aspect ratio has changed - axes = self.tab_active["axes"] + axes = self.tab_active.axes num_subplots = len(axes) figure_extent = self.figure.get_window_extent() figure_ratio = figure_extent.width / figure_extent.height @@ -232,7 +274,7 @@ def __click(self, event: Event, force_update: bool = False) -> None: # Get tab name to activate for tab in self.tabs_data.values(): - button = tab["button"] + button = tab.button if button.ax == event.inaxes: tab_name = button.label.get_text().replace('\n', ' ') break @@ -247,20 +289,20 @@ def __click(self, event: Event, force_update: bool = False) -> None: cur_stack = self.figure.canvas.toolbar._nav_stack for tab in self.tabs_data.values(): if self.sync_tabs or tab is self.tab_active: - tab["nav_stack"] = cur_stack._elements.copy() - tab["nav_pos"] = cur_stack._pos + tab.nav_stack = cur_stack._elements.copy() + tab.nav_pos = cur_stack._pos # Update axes and title - for ax in self.tab_active["axes"]: + for ax in self.tab_active.axes: self.subfigs[0].delaxes(ax) if self.legend is not None: self.legend.remove() self.legend = None self.tab_active = self.tabs_data[tab_name] self.subfigs[0].suptitle(tab_name) - for ax in self.tab_active["axes"]: + for ax in self.tab_active.axes: self.subfigs[0].add_subplot(ax) - handles, labels = self.tab_active["legend_data"] + handles, labels = self.tab_active.legend_data if labels: self.legend = self.subfigs[0].legend( handles, labels, ncol=len(handles), loc='outside lower center') @@ -268,13 +310,13 @@ def __click(self, event: Event, force_update: bool = False) -> None: # # Restore navigation history and toolbar state if necessary if not self.offscreen: assert self.figure.canvas.toolbar is not None - cur_stack._elements = self.tab_active["nav_stack"] - cur_stack._pos = self.tab_active["nav_pos"] + cur_stack._elements = self.tab_active.nav_stack + cur_stack._pos = self.tab_active.nav_pos self.figure.canvas.toolbar.set_history_buttons() # Update buttons style for tab in self.tabs_data.values(): - button = tab["button"] + button = tab.button if tab is self.tab_active: button.ax.set_facecolor('green') button.color = 'green' @@ -291,7 +333,8 @@ def refresh(self) -> None: """Refresh canvas drawing. """ self.figure.canvas.draw() - self.figure.canvas.flush_events() + # FIXME: `flush_events` on 'matplotlib>=3.8' causes deadlock + # self.figure.canvas.flush_events() def add_tab(self, # pylint: disable=unused-argument tab_name: str, @@ -338,7 +381,7 @@ def add_tab(self, # pylint: disable=unused-argument Optional: True by default. """ # Make sure that the time sequence is valid - assert (np.diff(time) > 0.0).all(), ( + assert (np.diff(time) >= 0.0).all(), ( "The time sequence must be monotonically increasing.") # Make sure that the provided tab name does not exist already @@ -420,14 +463,8 @@ def add_tab(self, # pylint: disable=unused-argument button.on_clicked(self.__click) # Create new tab data container - self.tabs_data[tab_name] = { - "axes": axes, - "legend_data": legend_data, - "nav_stack": [], - "nav_pos": -1, - "button": button, - "button_axcut": button_axcut - } + self.tabs_data[tab_name] = TabData( + axes, legend_data, button, button_axcut, nav_stack=[], nav_pos=-1) # Check if it is the first tab to be added if self.tab_active is None: @@ -465,7 +502,7 @@ def select_active_tab(self, tab_name: str) -> None: "No tab with this exact name has been added.") event = LocationEvent("click", self.figure.canvas, 0, 0) - event.inaxes = self.tabs_data[tab_name]["button"].ax + event.inaxes = self.tabs_data[tab_name].button.ax self.__click(event, force_update=True) def remove_tab(self, @@ -494,22 +531,22 @@ def remove_tab(self, self.select_active_tab(next(iter(self.tabs_data.keys()))) # Change reference axis if to be deleted - if any(ax is self.ref_ax for ax in tab["axes"]): + if any(ax is self.ref_ax for ax in tab.axes): if self.tabs_data: - self.ref_ax = self.tab_active["axes"][0] + self.ref_ax = self.tab_active.axes[0] else: self.ref_ax = None # Disable button - tab["button"].disconnect_events() - tab["button_axcut"].remove() + tab.button.disconnect_events() + tab.button_axcut.remove() # Remove axes and legend manually is not more tabs available if not self.tabs_data: if self.subfigs[0]._suptitle is not None: self.subfigs[0]._suptitle.remove() self.subfigs[0]._suptitle = None - for ax in tab["axes"]: + for ax in tab.axes: ax.remove() if self.legend is not None: self.legend.remove() @@ -632,9 +669,9 @@ def plot_log(log_data: Dict[str, Any], str, Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]] ] = OrderedDict() - # Get time and robot positions, velocities, and acceleration + # Get time and robot positions, velocities, accelerations and efforts time = log_vars["Global.Time"] - for fields_type in ("Position", "Velocity", "Acceleration"): + for fields_type in ("Position", "Velocity", "Acceleration", "Effort"): fieldnames: List[str] = getattr(robot, "_".join(( "log", fields_type.lower(), "fieldnames"))) if not enable_flexiblity_data: @@ -650,33 +687,25 @@ def plot_log(log_data: Dict[str, Any], tabs_data[' '.join(("State", fields_type))] = OrderedDict( (field[7:].replace(fields_type, ""), elem) for field, elem in values.items()) - except ValueError: + except KeyError: # Variable has not been recorded and is missing in log file pass - # Get motors efforts information - try: - motors_efforts = extract_variables_from_log( - log_vars, robot.log_motor_effort_fieldnames, namespace=robot.name) - tabs_data['MotorEffort'] = OrderedDict(zip( - robot.motor_names, motors_efforts)) - except ValueError: - # Variable has not been recorded and is missing in log file - pass - # Get command information try: command = extract_variables_from_log( log_vars, robot.log_command_fieldnames, namespace=robot.name) - tabs_data['Command'] = OrderedDict(zip(robot.motor_names, command)) - except ValueError: + tabs_data['Command'] = OrderedDict( + zip((motor.name for motor in robot.motors), command)) + except KeyError: # Variable has not been recorded and is missing in log file pass # Get sensors information for sensors_class, sensors_fields in SENSORS_FIELDS.items(): sensors_type: str = cast(str, sensors_class.type) - sensor_names = robot.sensor_names.get(sensors_type, []) + sensor_names = tuple( + sensor.name for sensor in robot.sensors.get(sensors_type, [])) if not sensor_names: continue if isinstance(sensors_fields, dict): @@ -691,7 +720,7 @@ def plot_log(log_data: Dict[str, Any], tabs_data[type_name] = OrderedDict( (field, OrderedDict(zip(sensor_names, data))) for field, data in zip(fieldnames, data_nested)) - except ValueError: + except KeyError: # Variable has not been recorded and is missing in log file pass else: @@ -703,7 +732,7 @@ def plot_log(log_data: Dict[str, Any], for name in sensor_names], robot.name) tabs_data[type_name] = OrderedDict(zip( sensor_names, data)) - except ValueError: + except KeyError: # Variable has not been recorded and is missing in log file pass @@ -763,7 +792,8 @@ def plot_log_interactive() -> None: # Load log file main_fullpath = main_arguments.input - log_vars = read_log(main_fullpath)["variables"] + log_data = read_log(main_fullpath) + log_vars = log_data["variables"] # If no plotting commands, display the list of headers instead if len(plotting_commands) == 0: @@ -778,7 +808,8 @@ def plot_log_interactive() -> None: if fullpath == main_fullpath or fullpath in compare_data.keys(): raise RuntimeError( "All log files must be unique when comparing them.") - compare_data[fullpath] = read_log(fullpath)["variables"] + log_data = read_log(fullpath) + compare_data[fullpath] = log_data["variables"] # Define line style cycle used for logs comparison linestyles = ("--", "-.", ":") diff --git a/python/jiminy_py/src/jiminy_py/robot.py b/python/jiminy_py/src/jiminy_py/robot.py index efa615871..d07273889 100644 --- a/python/jiminy_py/src/jiminy_py/robot.py +++ b/python/jiminy_py/src/jiminy_py/robot.py @@ -24,16 +24,11 @@ from . import core as jiminy from .core import ( # pylint: disable=no-name-in-module - EncoderSensor as encoder, - EffortSensor as effort, - ContactSensor as contact, - ForceSensor as force, - ImuSensor as imu) + SimpleMotor, EncoderSensor, EffortSensor, ForceSensor, ImuSensor) DEFAULT_UPDATE_RATE = 1000.0 # [Hz] DEFAULT_FRICTION_DRY_SLOPE = 20.0 -DEFAULT_ARMATURE = 0.1 EXTENSION_MODULES: Sequence[ModuleType] = () @@ -206,6 +201,8 @@ def generate_default_hardware_description_file( Motor=defaultdict(OrderedDict), Sensor=defaultdict(OrderedDict) ) + motors_info = hardware_info['Motor'] + sensors_info = hardware_info['Sensor'] # Extract the root link. It is the one having no parent at all. links = set() @@ -261,10 +258,10 @@ def generate_default_hardware_description_file( # Extract the sensor type sensor_type = gazebo_sensor_descr.attrib['type'].casefold() if 'imu' in sensor_type: - sensor_type = imu.type + sensor_type = ImuSensor.type elif 'contact' in sensor_type: collision_body_names.add(body_name) - sensor_type = force.type + sensor_type = ForceSensor.type else: LOGGER.warning( "Unsupported Gazebo sensor plugin of type '%s'.", @@ -296,7 +293,7 @@ def generate_default_hardware_description_file( map(float, frame_pose_obj.text.split())) # Add the sensor to the robot's hardware - hardware_info['Sensor'][sensor_type][sensor_name] = sensor_info + sensors_info[sensor_type][sensor_name] = sensor_info # Extract the collision bodies and ground model, then add force # sensors. At this point, every force sensor is associated with a @@ -304,7 +301,7 @@ def generate_default_hardware_description_file( if gazebo_plugin_descr.find('kp') is not None: # Add a force sensor, if not already in the collision set if body_name not in collision_body_names: - force_sensor_info = hardware_info['Sensor'][force.type] + force_sensor_info = sensors_info[ForceSensor.type] force_sensor_info[f"{body_name}Contact"] = OrderedDict( frame_name=body_name) @@ -335,22 +332,22 @@ def generate_default_hardware_description_file( body_name_obj = gazebo_plugin_descr.find('bodyName') assert body_name_obj is not None body_name = body_name_obj.text - force_sensor_info = hardware_info['Sensor'][force.type] + force_sensor_info = sensors_info[ForceSensor.type] force_sensor_info[f"{body_name}Wrench"] = OrderedDict( frame_name=body_name) else: LOGGER.warning("Unsupported Gazebo plugin '%s'", plugin) # Add IMU sensor to the root link if no Gazebo IMU sensor has been found - if link_root and imu.type not in hardware_info['Sensor'].keys(): - hardware_info['Sensor'][imu.type][link_root] = OrderedDict( + if link_root and ImuSensor.type not in sensors_info.keys(): + sensors_info[ImuSensor.type][link_root] = OrderedDict( frame_name=link_root) # Add force sensors and collision bodies if no Gazebo plugin is available if not gazebo_plugins_found: for link_leaf in links_leaf: # Add a force sensor - hardware_info['Sensor'][force.type][link_leaf] = OrderedDict( + sensors_info[ForceSensor.type][link_leaf] = OrderedDict( frame_name=link_leaf) # Add the related body to the collision set if possible @@ -393,7 +390,7 @@ def generate_default_hardware_description_file( # Extract the motors and effort sensors. # It is done by reading 'transmission' field, that is part of # URDF standard, so it should be available on any URDF file. - transmission_found = root.find('transmission') is not None + joint_transmissions = set() for transmission_descr in root.iterfind('transmission'): # Assert(s) for type checker assert isinstance(transmission_descr, ET.Element) @@ -429,6 +426,7 @@ def generate_default_hardware_description_file( assert isinstance(joint_descr, ET.Element) joint_name = joint_descr.attrib['name'] motor_info['joint_name'] = joint_name + joint_transmissions.add(joint_name) # Make sure that the joint is revolute joint = root.find(f"./joint[@name='{joint_name}']") @@ -456,15 +454,15 @@ def generate_default_hardware_description_file( else: armature_txt = armature.text assert armature_txt is not None - motor_info['armature'] = float(armature_txt) * ( - motor_info['mechanicalReduction'] ** 2) + motor_info['armature'] = float(armature_txt) # Add dynamics property to motor info, if any motor_info.update(joints_options.pop(joint_name)) - # Add the motor and sensor to the robot's hardware - hardware_info['Motor']['SimpleMotor'][motor_name] = motor_info - hardware_info['Sensor'][effort.type][joint_name] = sensor_info + # Add the motor and sensors to the robot's hardware + motors_info[SimpleMotor.__name__][motor_name] = motor_info + sensors_info[EncoderSensor.type][motor_name] = sensor_info + sensors_info[EffortSensor.type][motor_name] = sensor_info # Define default encoder sensors, and default effort sensors if no # transmission available. @@ -481,20 +479,21 @@ def generate_default_hardware_description_file( encoder_info['joint_name'] = joint_name # Add the sensor to the robot's hardware - hardware_info['Sensor'][encoder.type][joint_name] = encoder_info + if joint_name not in joint_transmissions: + sensors_info[EncoderSensor.type][joint_name] = encoder_info # Add motors to robot hardware by default if no transmission found - if not transmission_found: + if not joint_transmissions: joint_limit_descr = joint_descr.find('./limit') assert joint_limit_descr is not None if float(joint_limit_descr.attrib['effort']) == 0.0: continue - hardware_info['Motor']['SimpleMotor'][joint_name] = OrderedDict( + motors_info[SimpleMotor.type][joint_name] = OrderedDict( joint_name=joint_name, - armature=DEFAULT_ARMATURE, + armature=0.0, **joints_options.pop(joint_name) ) - hardware_info['Sensor'][effort.type][joint_name] = OrderedDict( + sensors_info[EffortSensor.type][joint_name] = OrderedDict( motor_name=joint_name) # Warn if friction model has been defined for non-actuated joints @@ -755,52 +754,42 @@ def load_hardware_description_file( motor.set_options(options) # Add the sensors to the robot + motor_names = [motor.name for motor in robot.motors] for sensor_type, sensors_descr in sensors_info.items(): for sensor_name, sensor_descr in sensors_descr.items(): + # Extract initialization arguments from options + init_kwargs = { + field: sensor_descr.pop(field) + for field in ( + 'joint_name', + 'motor_name', + 'frame_name', + 'body_name', + 'frame_pose') + if field in sensor_descr} + # Make sure the sensor can be instantiated - if sensor_type == encoder.type: - joint_name = sensor_descr.pop('joint_name') + joint_name = init_kwargs.get('joint_name', None) + if joint_name is not None: + init_kwargs['joint_name'] = joint_name if not robot.pinocchio_model.existJointName(joint_name): - LOGGER.warning( - "'%s' is not a valid joint name.", joint_name) - continue - elif sensor_type == effort.type: - motor_name = sensor_descr.pop('motor_name') - if motor_name not in robot.motor_names: - LOGGER.warning( - "'%s' is not a valid motor name.", motor_name) - continue - - # Create the sensor and attach it - for module in (jiminy, *EXTENSION_MODULES): - try: - sensor = getattr(module, sensor_type)(sensor_name) - break - except AttributeError: - pass - else: - raise RuntimeError( - f"Cannot instantiate sensor of type '{sensor_type}'.") - robot.attach_sensor(sensor) + raise ValueError( + f"'{joint_name}' is not a valid joint name.") - # Initialize the sensor - if sensor_type == encoder.type: - sensor.initialize(joint_name) - elif sensor_type == effort.type: - sensor.initialize(motor_name) - elif sensor_type == contact.type: - frame_name = sensor_descr.pop('frame_name') - sensor.initialize(frame_name) - elif sensor_type in [force.type, imu.type]: - # Create the frame and add it to the robot model - frame_name = sensor_descr.pop('frame_name', None) + motor_name = init_kwargs.get('motor_name', None) + if motor_name is not None: + init_kwargs['motor_name'] = motor_name + if motor_name not in motor_names: + raise ValueError( + f"'{motor_name}' is not a valid motor name.") + frame_name = init_kwargs.get('frame_name', None) + if frame_name is not None: # Create a frame if a frame name has been specified. # In such a case, the body name must be specified. - if not frame_name or \ - not robot.pinocchio_model.existFrame(frame_name): + if not robot.pinocchio_model.existFrame(frame_name): # Get the body name - body_name = sensor_descr.pop('body_name') + body_name = init_kwargs.pop('body_name') # Generate a frame name both intelligible and available if frame_name is None: @@ -813,25 +802,36 @@ def load_hardware_description_file( i += 1 # Compute SE3 object representing the frame placement - frame_pose_xyzrpy = np.array( - sensor_descr.pop('frame_pose')) + frame_pose_xyzrpy = np.array(init_kwargs.pop('frame_pose')) frame_trans = frame_pose_xyzrpy[:3] frame_rot = rpyToMatrix(frame_pose_xyzrpy[3:]) frame_placement = pin.SE3(frame_rot, frame_trans) # Add the frame to the robot model robot.add_frame(frame_name, body_name, frame_placement) - elif 'frame_pose' in sensor_descr.keys(): + + # Add newly created frame to initialization options + init_kwargs['frame_name'] = frame_name + elif 'frame_pose' in init_kwargs.keys(): raise ValueError( f"The sensor '{sensor_name}' is attached to the frame " f"'{frame_name}' that already exists whereas a " "specific pose is also requested.") - # Initialize the sensor - sensor.initialize(frame_name) + # Create the sensor and attach it + for module in (jiminy, *EXTENSION_MODULES): + try: + sensor = getattr(module, sensor_type)(sensor_name) + break + except AttributeError: + pass else: - raise ValueError( - f"Unsupported sensor type {sensor_type}.") + raise RuntimeError( + f"Cannot instantiate sensor of type '{sensor_type}'.") + robot.attach_sensor(sensor) + + # Initialize the sensor + sensor.initialize(**init_kwargs) # Set the sensor options options = sensor.get_options() diff --git a/python/jiminy_py/src/jiminy_py/simulator.py b/python/jiminy_py/src/jiminy_py/simulator.py index 110912d60..06d3a2b43 100644 --- a/python/jiminy_py/src/jiminy_py/simulator.py +++ b/python/jiminy_py/src/jiminy_py/simulator.py @@ -20,7 +20,7 @@ from . import core as jiminy from .robot import BaseJiminyRobot, generate_default_hardware_description_file -from .dynamics import TrajectoryDataType +from .dynamics import Trajectory from .log import read_log, build_robot_from_log from .viewer import (CameraPoseType, interactive_mode, @@ -819,7 +819,7 @@ def render(self, def replay(self, extra_logs_files: Sequence[str] = (), - extra_trajectories: Sequence[TrajectoryDataType] = (), + extra_trajectories: Sequence[Trajectory] = (), **kwargs: Any) -> None: """Replay the current episode until now. @@ -850,7 +850,7 @@ def replay(self, logs_data.append(log_data) # Extract trajectory data from pairs (robot, log) - trajectories: List[TrajectoryDataType] = [] + trajectories: List[Trajectory] = [] update_hooks: List[ Optional[Callable[[float, np.ndarray, np.ndarray], None]]] = [] extra_kwargs: Dict[str, Any] = {} diff --git a/python/jiminy_py/src/jiminy_py/viewer/meshcat/wrapper.py b/python/jiminy_py/src/jiminy_py/viewer/meshcat/wrapper.py index aa600e091..edbda3f83 100644 --- a/python/jiminy_py/src/jiminy_py/viewer/meshcat/wrapper.py +++ b/python/jiminy_py/src/jiminy_py/viewer/meshcat/wrapper.py @@ -100,7 +100,7 @@ def __call__(self) -> None: msg = None else: # Do not spend time analyzing messages already rejected - msg = None + idents, msg = None, None if msg is not None and \ msg['header']['msg_type'].startswith('comm_'): diff --git a/python/jiminy_py/src/jiminy_py/viewer/replay.py b/python/jiminy_py/src/jiminy_py/viewer/replay.py index 9a2d34b78..fc7e9add3 100644 --- a/python/jiminy_py/src/jiminy_py/viewer/replay.py +++ b/python/jiminy_py/src/jiminy_py/viewer/replay.py @@ -11,7 +11,6 @@ import tempfile import argparse from base64 import b64encode -from bisect import bisect_right from collections import deque from types import TracebackType from functools import wraps, partial @@ -28,7 +27,7 @@ import pinocchio as pin from .. import core as jiminy -from ..dynamics import TrajectoryDataType +from ..dynamics import Trajectory from ..log import (read_log, build_robot_from_log, extract_trajectory_from_log, @@ -38,7 +37,6 @@ Tuple4FType, CameraPoseType, CameraMotionType, - interp1d, get_default_backend, is_display_available, Viewer) @@ -150,8 +148,8 @@ def fun_safe(*args: Any, **kwargs: Any) -> Any: @_with_lock def play_trajectories( - trajs_data: Union[ # pylint: disable=unused-argument - TrajectoryDataType, Sequence[TrajectoryDataType]], + trajectories: Union[ # pylint: disable=unused-argument + Trajectory, Sequence[Trajectory]], update_hooks: Optional[Union[ Callable[[float, np.ndarray, np.ndarray], None], Sequence[Optional[Callable[[float, np.ndarray, np.ndarray], None]]] @@ -193,7 +191,7 @@ def play_trajectories( Replay speed is independent of the platform (windows, linux...) and available CPU power. - :param trajs_data: List of `TrajectoryDataType` dicts. + :param trajectories: List of trajectories. :param update_hooks: Callables associated with each robot that can be used to update non-kinematic robot data, for instance to emulate sensors data from log using the hook provided @@ -226,7 +224,7 @@ def play_trajectories( Optional: None by default. :param enable_travelling: Whether the camera tracks the robot associated with the first trajectory specified in - `trajs_data`. `None` to disable. + `trajectories`. `None` to disable. Optional: Disabled by default. :param camera_motion: Camera breakpoint poses over time, as a list of `CameraMotionBreakpointType` dict. None to disable. @@ -308,10 +306,11 @@ def play_trajectories( :returns: List of viewers used to play the trajectories. """ # Make sure sequence arguments are list or tuple - if isinstance(trajs_data, dict): - trajs_data = [trajs_data] + if isinstance(trajectories, Trajectory): + trajectories = [trajectories] + ntrajs = len(trajectories) if update_hooks is None: - update_hooks = [None] * len(trajs_data) + update_hooks = [None] * ntrajs if callable(update_hooks): update_hooks = [update_hooks] if isinstance(viewers, Viewer): @@ -379,8 +378,8 @@ def play_trajectories( if viewers is None and backend.startswith('panda3d'): # Check whether at least one of the robots has a freeflyer has_freeflyer = False - for traj in trajs_data: - robot = traj['robot'] + for trajectory in trajectories: + robot = trajectory.robot assert robot is not None if robot.has_freeflyer: has_freeflyer = True @@ -396,8 +395,8 @@ def play_trajectories( # Make sure it is possible to display contacts if requested if display_contacts: - for traj in trajs_data: - robot = traj['robot'] + for trajectory in trajectories: + robot = trajectory.robot assert robot is not None if robot.is_locked: LOGGER.debug( @@ -408,23 +407,23 @@ def play_trajectories( # Sanitize user-specified robot offsets if xyz_offsets is None: - xyz_offsets = len(trajs_data) * (None,) - elif len(xyz_offsets) != len(trajs_data): + xyz_offsets = ntrajs * (None,) + elif len(xyz_offsets) != ntrajs: assert isinstance(xyz_offsets[0], float) xyz_offsets = np.tile( - xyz_offsets, (len(trajs_data), 1)) # type: ignore[arg-type] + xyz_offsets, (ntrajs, 1)) # type: ignore[arg-type] assert xyz_offsets is not None # Sanitize user-specified robot colors if robots_colors is None: - if len(trajs_data) == 1: + if ntrajs == 1: robots_colors = (None,) else: robots_colors = list(islice( - cycle(COLORS.values()), len(trajs_data))) + cycle(COLORS.values()), ntrajs)) elif isinstance(robots_colors, str) or isinstance(robots_colors[0], float): robots_colors = [robots_colors] # type: ignore[list-item] - assert len(robots_colors) == len(trajs_data) + assert len(robots_colors) == ntrajs # Sanitize user-specified legend if legend is not None and isinstance(legend, str): @@ -432,21 +431,21 @@ def play_trajectories( # Make sure the viewers instances are consistent with the trajectories if viewers is None: - viewers = [None for _ in trajs_data] - assert len(viewers) == len(trajs_data) + viewers = [None for _ in trajectories] + assert len(viewers) == ntrajs # Instantiate or refresh viewers if necessary - for i, (viewer, traj, color) in enumerate(zip( - viewers, trajs_data, robots_colors)): + for i, (viewer, trajectory, color) in enumerate(zip( + viewers, trajectories, robots_colors)): # Extract robot from trajectory - robot = traj['robot'] + robot = trajectory.robot assert robot is not None # Create new viewer instance if necessary, and load the robot in it if viewer is None: uniq_id = next(tempfile._get_candidate_names()) robot_name = f"{uniq_id}_robot_{i}" - use_theoretical_model = traj['use_theoretical_model'] + use_theoretical_model = trajectory.use_theoretical_model viewer = Viewer( robot, use_theoretical_model=use_theoretical_model, @@ -459,7 +458,7 @@ def play_trajectories( viewers[i] = viewer # Reset the configuration of the robot - if traj['use_theoretical_model']: + if trajectory.use_theoretical_model: model = robot.pinocchio_model_th else: model = robot.pinocchio_model @@ -474,15 +473,12 @@ def play_trajectories( # Assert(s) for type checker assert all(viewers) viewers = cast(List[Viewer], viewers) + assert Viewer._backend_obj is not None # Add default legend with robots names if replaying multiple trajectories if all(color is not None for color in robots_colors) and legend is None: legend = [viewer.robot_name for viewer in viewers] - # Use first viewers as main viewer to call static methods conveniently - viewer = viewers[0] - assert Viewer._backend_obj is not None - # Make sure clock is only enabled for panda3d backend if enable_clock and not backend.startswith('panda3d'): LOGGER.warning( @@ -490,7 +486,7 @@ def play_trajectories( enable_clock = False # Early return if nothing to replay - if all(not traj['evolution_robot'] for traj in trajs_data): + if all(not trajectory.states for trajectory in trajectories): LOGGER.debug("Nothing to replay.") return viewers @@ -515,33 +511,34 @@ def play_trajectories( if time_interval[1] < time_interval[0]: raise ValueError("Time interval must be non-empty and positive.") - # Initialize robot configuration is viewer before any further processing - for viewer_i, traj, offset in zip(viewers, trajs_data, xyz_offsets): - data = traj['evolution_robot'] - if data: - i = bisect_right( - [s.t for s in data], time_interval[0], hi=len(data)-1) - for f_ext in viewer_i.f_external: - f_ext.vector[:] = 0.0 - viewer_i.display(data[i].q, data[i].v, offset) + # Initialize robot and viewer configuration before any further processing + for viewer, trajectory, offset in zip(viewers, trajectories, xyz_offsets): + try: + state = trajectory.get(time_interval[0], mode='clip') + except RuntimeError: + pass + else: + viewer.display(state.q, state.v, offset) + for f_ext in viewer.f_external: + f_ext.vector[:] = 0.0 if backend.startswith('panda3d'): if display_com is not None: - viewer_i.display_center_of_mass(display_com) + viewer.display_center_of_mass(display_com) if display_dcm is not None: - viewer_i.display_capture_point(display_dcm) + viewer.display_capture_point(display_dcm) if display_contacts is not None: - viewer_i.display_contact_forces(display_contacts) + viewer.display_contact_forces(display_contacts) if display_f_external is not None: - viewer_i.display_external_forces(display_f_external) + viewer.display_external_forces(display_f_external) # Set camera pose or activate camera travelling if requested + viewer, robot = viewers[0], trajectories[0].robot if enable_travelling: position, rotation, relative = None, None, None if camera_pose is not None: position, rotation, relative = camera_pose if relative is None: # Track the first actual frame by default (0: world, 1: root_joint) - robot = trajs_data[0]['robot'] assert robot is not None if not robot.has_freeflyer: raise ValueError( @@ -567,57 +564,12 @@ def play_trajectories( if record_video_path is not None: # Extract and resample trajectory data at fixed framerate time_max = time_interval[0] - for traj in trajs_data: - if len(traj['evolution_robot']): - time_max = max([time_max, traj['evolution_robot'][-1].t]) + for trajectory in trajectories: + if trajectory.has_data: + _, t_end = trajectory.time_interval + time_max = max(time_max, t_end) time_max = min(time_max, time_interval[1]) - time_global = np.arange( - time_interval[0], time_max, speed_ratio / VIDEO_FRAMERATE) - position_evolutions: List[Optional[np.ndarray]] = [] - velocity_evolutions: List[Optional[ - Union[np.ndarray, Sequence[None]]]] = [] - force_evolutions: List[Optional[ - Union[List[List[np.ndarray]], Sequence[None]]]] = [] - for traj in trajs_data: - if len(traj['evolution_robot']): - data_orig = traj['evolution_robot'] - robot = traj['robot'] - assert robot is not None - if traj['use_theoretical_model']: - model = robot.pinocchio_model_th - else: - model = robot.pinocchio_model - t_orig = np.array([s.t for s in data_orig]) - pos_orig = np.stack([s.q for s in data_orig], axis=0) - position_evolutions.append(jiminy.interpolate_positions( - model, t_orig, pos_orig.T, time_global).T) - if data_orig[0].v is not None: - vel_orig = np.stack([ - s.v # type: ignore[misc] - for s in data_orig], axis=0) - velocity_evolutions.append(interp1d( - t_orig, vel_orig, time_global)) - else: - velocity_evolutions.append((None,) * len(time_global)) - if data_orig[0].f_ext is not None: - forces: List[np.ndarray] = [] - for i in range(len(data_orig[0].f_ext)): - f_ext_orig = np.stack([ - s.f_ext[i] # type: ignore[index] - for s in data_orig], axis=0) - forces.append(interp1d( - t_orig, f_ext_orig, time_global)) - force_evolutions.append([ - [f_ext[i] for f_ext in forces] - for i in range(len(time_global))]) - else: - force_evolutions.append((None,) * len(time_global)) - else: - position_evolutions.append(None) - velocity_evolutions.append(None) - force_evolutions.append(None) - # Initialize video recording if backend == 'meshcat': # Sanitize the recording path to enforce '.webm' extension @@ -645,30 +597,37 @@ def play_trajectories( frame = av.VideoFrame(*record_video_size, 'rgb24') # Add frames to video sequentially - for i, t_cur in enumerate(tqdm( + update_hook_t = None + time_global = np.arange( + time_interval[0], time_max + 1e-10, speed_ratio / VIDEO_FRAMERATE) + for t in tqdm( time_global, desc="Rendering frames", - disable=(not verbose and not record_video_html_embedded))): + disable=(not verbose and not record_video_html_embedded)): try: - # Update 3D view - for viewer, pos, vel, forces, xyz_offset, update_hook in zip( - viewers, position_evolutions, velocity_evolutions, - force_evolutions, xyz_offsets, update_hooks): + # Loop over all trajectories + for viewer, trajectory, xyz_offset, update_hook in zip( + viewers, trajectories, xyz_offsets, update_hooks): + # Skip empty trajectories assert viewer is not None - if pos is None: + if not trajectory.has_data: continue - q, v, f_ext = pos[i], vel[i], forces[i] - if f_ext is not None: - for f_ref, f_i in zip(viewer.f_external, f_ext): + + # Compute robot state at current time + state = trajectory.get(t, mode='clip') + + # Update viewer state + if state.f_external is not None: + for f_ref, f_i in zip( + viewer.f_external, state.f_external[1:]): f_ref.vector[:] = f_i if update_hook is not None: - update_hook_t = partial(update_hook, t_cur, q, v) - else: - update_hook_t = None - viewer.display(q, v, xyz_offset, update_hook_t) + update_hook_t = partial( + update_hook, t, state.q, state.v) + viewer.display(state.q, state.v, xyz_offset, update_hook_t) # Update clock if enabled if enable_clock: - Viewer.set_clock(t_cur) + Viewer.set_clock(t) # Add frame to video if backend == 'meshcat': @@ -721,12 +680,12 @@ def replay_thread(viewer: Viewer, *args: Any) -> None: viewer.replay(*args) threads = [] - for viewer, traj, xyz_offset, update_hook in zip( - viewers, trajs_data, xyz_offsets, update_hooks): + for viewer, trajectory, xyz_offset, update_hook in zip( + viewers, trajectories, xyz_offsets, update_hooks): threads.append(Thread( target=replay_thread, args=(viewer, - traj['evolution_robot'], + trajectory, time_interval, speed_ratio, xyz_offset, @@ -785,8 +744,8 @@ def replay_thread(viewer: Viewer, *args: Any) -> None: def extract_replay_data_from_log( log_data: Dict[str, np.ndarray], - robot: jiminy.Robot) -> Tuple[ - TrajectoryDataType, + robot: Optional[jiminy.Robot] = None) -> Tuple[ + Trajectory, Optional[Callable[[float, np.ndarray, np.ndarray], None]], Dict[str, Any]]: """Extract replay data from log data. @@ -799,14 +758,14 @@ def extract_replay_data_from_log( By default, it enables display of external forces applied on freeflyer if any. """ - # For each pair (log, robot), extract a trajectory object for - # `play_trajectories` + # Extract a trajectory object for `play_trajectories` trajectory = extract_trajectory_from_log(log_data, robot) + robot = trajectory.robot # Display external forces on root joint, if any replay_kwargs = {} if robot.has_freeflyer: - if trajectory['use_theoretical_model']: + if trajectory.use_theoretical_model: njoints = robot.pinocchio_model_th.njoints else: njoints = robot.pinocchio_model.njoints @@ -817,7 +776,7 @@ def extract_replay_data_from_log( if not robot.is_locked: update_hook = update_sensor_measurements_from_log(log_data, robot) else: - if robot.sensor_names: + if robot.sensors: LOGGER.warning( "At least one of the robot is locked, which means that a " "simulation using the robot is still running. It will be " @@ -828,17 +787,19 @@ def extract_replay_data_from_log( return trajectory, update_hook, replay_kwargs -def play_logs_data(robots: Union[Sequence[jiminy.Robot], jiminy.Robot], - logs_data: Union[Sequence[Dict[str, np.ndarray]], - Dict[str, np.ndarray]], - **kwargs: Any) -> Sequence[Viewer]: +def play_logs_data( + logs_data: Union[ + Sequence[Dict[str, np.ndarray]], Dict[str, np.ndarray]], + robots: Optional[Union[ + Sequence[Optional[jiminy.Robot]], jiminy.Robot]] = None, + **kwargs: Any) -> Sequence[Viewer]: """Play log data in a viewer. This method simply formats the data then calls `play_trajectories`. - :param robots: Either a single robot, or a list of robot for each log data. :param logs_data: Either a single dictionary, or a list of dictionaries of simulation data log. + :param robots: Either a single robot, or a list of robot for each log data. :param kwargs: Keyword arguments to forward to `play_trajectories` method. """ # Reformat input arguments as lists @@ -846,17 +807,19 @@ def play_logs_data(robots: Union[Sequence[jiminy.Robot], jiminy.Robot], logs_data = [logs_data] if isinstance(robots, jiminy.Robot): robots = [robots] + elif robots is None: + robots = [None,] * len(logs_data) # Extract a replay data for `play_trajectories` for each pair (robot, log) trajectories, update_hooks, extra_kwargs = [], [], {} - for robot, log_data in zip(robots, logs_data): - traj, update_hook, _kwargs = \ + for log_data, robot in zip(logs_data, robots): + trajectory, update_hook, _kwargs = \ extract_replay_data_from_log(log_data, robot) - trajectories.append(traj) + trajectories.append(trajectory) update_hooks.append(update_hook) extra_kwargs.update(_kwargs) - # Do not display external forces by default if replaying several traj + # Do not display external forces by default if replaying several trajectory if len(trajectories) > 1: extra_kwargs.pop("display_f_external", None) @@ -904,7 +867,7 @@ def play_logs_files(logs_files: Union[str, Sequence[str]], for log_file in logs_files] # Forward arguments to lower-level method - return play_logs_data(robots, logs_data, **kwargs) + return play_logs_data(logs_data, robots, **kwargs) def async_play_and_record_logs_files( @@ -1062,11 +1025,3 @@ def _play_logs_files_entrypoint() -> None: # Do not exit method as long as a graphical window is open while Viewer.has_gui(): time.sleep(0.5) - - -__all__ = [ - 'extract_replay_data_from_log', - 'play_logs_data', - 'play_logs_files', - 'async_play_and_record_logs_files', -] diff --git a/python/jiminy_py/src/jiminy_py/viewer/viewer.py b/python/jiminy_py/src/jiminy_py/viewer/viewer.py index 82e5b3ef3..4785ed12f 100644 --- a/python/jiminy_py/src/jiminy_py/viewer/viewer.py +++ b/python/jiminy_py/src/jiminy_py/viewer/viewer.py @@ -1,5 +1,10 @@ # mypy: disable-error-code="attr-defined, name-defined" -""" TODO: Write documentation. +"""This module provides a backend-agnostic 3D viewer on top of lower-level +implementations. + +This viewer can be used to render the state of a robot, including sensor data +and external forces if any. It is mainly used to render the current simulation +state or replay multiple complete trajectories as-posteriori. """ import os import re @@ -19,7 +24,6 @@ from copy import deepcopy from urllib.request import urlopen from functools import wraps, partial -from bisect import bisect_right from threading import RLock from multiprocessing import Process as ProcessMP from typing import ( @@ -42,10 +46,10 @@ from .. import core as jiminy from ..core import ( # pylint: disable=no-name-in-module - ContactSensor as contact, + ContactSensor, discretize_heightmap) from ..robot import _DuplicateFilter -from ..dynamics import State +from ..dynamics import Trajectory from .meshcat.utilities import interactive_mode from .panda3d.panda3d_visualizer import ( Tuple3FType, Tuple4FType, ShapeType, Panda3dApp, Panda3dViewer, @@ -253,7 +257,11 @@ def sleep(dt: float) -> None: def get_color_code( color: Optional[Union[str, Tuple4FType]]) -> Optional[Tuple4FType]: - """ TODO: Write documentation. + """Sanitize color codes. + + All it does is converting color codes that has been specified through + pre-defined names into their corresponding 4-tuple (R, G, B, A), where each + component is a floating-point scalar normalized in range [0.0, 1.0]. """ if isinstance(color, str): try: @@ -319,12 +327,13 @@ def __del__(self) -> None: self.kill() def is_parent(self) -> bool: - """ TODO: Write documentation. + """Whether the wrapped process is a child of the main python process. """ return not isinstance(self._proc, Process) def is_alive(self) -> bool: - """ TODO: Write documentation. + """Whether the wrapped process is running or idle, but not terminated + nor a zombie yet. """ if isinstance(self._proc, ProcessMP): return self._proc.is_alive() @@ -333,25 +342,27 @@ def is_alive(self) -> bool: if isinstance(self._proc, Process): import psutil # pylint: disable=import-outside-toplevel try: - return self._proc.status() in [ - psutil.STATUS_RUNNING, psutil.STATUS_SLEEPING] + return self._proc.status() in ( + psutil.STATUS_RUNNING, psutil.STATUS_SLEEPING) except psutil.NoSuchProcess: return False # if isinstance(self._proc, Panda3dApp): return hasattr(self._proc, 'win') - def wait(self, timeout: Optional[float] = None) -> None: - """ TODO: Write documentation. + def wait(self, timeout: float) -> None: + """Wait until the wrapped process finish what its current task or + timeout is reached. """ if isinstance(self._proc, ProcessMP): self._proc.join(timeout) if isinstance(self._proc, (subprocess.Popen, Process)): - self._proc.wait(timeout) # type: ignore[arg-type] + self._proc.wait(timeout) if isinstance(self._proc, Panda3dApp): self._proc.step() def kill(self) -> None: - """ TODO: Write documentation. + """If wrapped process if it is a child of the main python process and + it is still alive that the time being. """ if self.is_parent() and self.is_alive(): if isinstance(self._proc, Panda3dApp): @@ -386,13 +397,19 @@ def kill(self) -> None: class CameraMotionBreakpointType(TypedDict, total=True): - """Time + """Basic data structure storing all the information needed to describe a + camera motion breakpoint, ie the expected pose of the camera at a given + point in time. """ + t: float - """Absolute pose of the camera, as a tuple position [X, Y, Z], rotation - [Roll, Pitch, Yaw]. + """Time at which the camera is expected to have in a specific pose. """ + pose: Tuple[Tuple3FType, Tuple3FType] + """Exacted absolute pose of the camera, as a tuple position (X, Y, Z), + rotation (Roll, Pitch, Yaw). + """ CameraMotionType = Sequence[CameraMotionBreakpointType] @@ -418,7 +435,11 @@ class MarkerDataType(TypedDict, total=True): class Viewer: - """ TODO: Write documentation. + """Backend-agnostic 3D viewer. + + This class can be used to render the state of a robot, including sensor + data and external forces if any. It is mainly used to render the current + simulation state or replay multiple complete trajectories as-posteriori. .. note:: The environment variable 'JIMINY_INTERACTIVE_DISABLE' can be @@ -816,19 +837,17 @@ def get_contact_scale(sensor_data: np.ndarray) -> Tuple3FType: f_z_rel = sensor_data[2] / CONTACT_FORCE_SCALE return (1.0, 1.0, min(max(f_z_rel, -1.0), 1.0)) - if contact.type in robot.sensor_names.keys(): - for name in robot.sensor_names[contact.type]: - sensor = robot.get_sensor(contact.type, name) - frame_index, data = sensor.frame_index, sensor.data - self.add_marker(name='_'.join((contact.type, name)), - shape="cylinder", - pose=self._client.data.oMf[frame_index], - scale=partial(get_contact_scale, data), - remove_if_exists=True, - auto_refresh=False, - radius=0.02, - length=0.5, - anchor_bottom=True) + for sensor in robot.sensors.get(ContactSensor.type, []): + frame_index, data = sensor.frame_index, sensor.data + self.add_marker(name='_'.join((sensor.type, sensor.name)), + shape="cylinder", + pose=self._client.data.oMf[frame_index], + scale=partial(get_contact_scale, data), + remove_if_exists=True, + auto_refresh=False, + radius=0.02, + length=0.5, + anchor_bottom=True) self.display_contact_forces(self._display_contact_forces) @@ -1019,7 +1038,9 @@ def open_gui() -> None: @staticmethod def has_gui() -> bool: - """ TODO: Write documentation. + """Whether the viewer has a Graphical User Interface (GUI), ie the + viewer is connected to a given backend that is running up to now, and + onscreen rendering has been specifically requested. """ if Viewer.is_alive(): # Assert(s) for type checker @@ -1497,7 +1518,7 @@ def set_camera_transform(self: Optional["Viewer"] = None, name and index in model. :param wait: Whether to wait for rendering to finish. """ - # pylint: disable=invalid-name + # pylint: disable=invalid-name, possibly-used-before-assignment # Assert(s) for type checker assert Viewer.backend is not None assert Viewer._backend_obj is not None @@ -1887,6 +1908,7 @@ def capture_frame(width: Optional[int] = None, # Extract and return numpy array RGB return np.frombuffer(buffer, np.uint8).reshape((height, width, 3)) + # if Viewer.backend == 'meshcat': # Send capture frame request to the background recorder process img_html = Viewer._backend_obj.capture_frame(width, height) @@ -2208,7 +2230,7 @@ def display_contact_forces(self, visibility: bool) -> None: # Update visibility for name in self.markers: - if name.startswith(contact.type): + if name.startswith(ContactSensor.type): self._gui.show_node(self._markers_group, name, visibility) self._markers_visibility[name] = visibility self._display_contact_forces = visibility @@ -2480,7 +2502,7 @@ def display(self, @_must_be_open def replay(self, - evolution_robot: Sequence[State], + trajectory: Trajectory, time_interval: Union[ np.ndarray, Tuple[float, float]] = (0.0, np.inf), speed_ratio: float = 1.0, @@ -2500,7 +2522,7 @@ def replay(self, It will alter original robot data if viewer attribute `use_theoretical_model` is false. - :param evolution_robot: List of State object of increasing time. + :param states: Sequence of `State` objects of increasing time. :param time_interval: Specific time interval to replay. Optional: Complete evolution by default [0, inf]. :param speed_ratio: Real-time factor. @@ -2519,6 +2541,14 @@ def replay(self, Optional: No update hook by default. :param wait: Whether to wait for rendering to finish. """ + # Early return if nothing to replay + if not trajectory.has_data: + return + + # Sanitize replay time interval + t_start, t_end = trajectory.time_interval + t_start, t_end = (min(max(t, t_start), t_end) for t in time_interval) + # Disable display of sensor data if no update hook is provided disable_display_contact_forces = False if update_hook is None and self._display_contact_forces: @@ -2527,68 +2557,53 @@ def replay(self, # Disable display of DCM if no velocity data provided disable_display_dcm = False - has_velocities = evolution_robot[0].v is not None - if not has_velocities and self._display_dcm: + if not trajectory.has_velocity and self._display_dcm: disable_display_dcm = True self.display_capture_point(False) - # Check if force data is available - has_forces = evolution_robot[0].f_ext is not None - # Replay the whole trajectory at constant speed ratio - v = None - update_hook_t = None - times = [s.t for s in evolution_robot] - t_simu = time_interval[0] - i = bisect_right(times, t_simu) + t = t_start time_init = time.time() time_prev = time_init - while i < len(evolution_robot): + v, update_hook_t = None, None + while True: try: # Update clock if enabled if enable_clock: - Viewer.set_clock(t_simu) - - # Compute interpolated data at current time - s_next = evolution_robot[min(i, len(times) - 1)] - s = evolution_robot[max(i - 1, 0)] - ratio = (t_simu - s.t) / (s_next.t - s.t) - q = pin.interpolate(self._client.model, s.q, s_next.q, ratio) - if has_velocities: - v = s.v + ratio * ( - s_next.v - s.v) # type: ignore[operator] - if has_forces: - for i, (f_ext, f_ext_next) in enumerate(zip( - s.f_ext, s_next.f_ext)): # type: ignore[arg-type] - self.f_external[i].vector[:] = \ - f_ext + ratio * (f_ext_next - f_ext) + Viewer.set_clock(t) + + # Compute state at current time + state = trajectory.get(t, mode="raise") + + # Update viewer force buffer + if state.f_external is not None: + for f_ref, f_i in zip( + self.f_external, state.f_external[1:]): + f_ref.vector[:] = f_i # Update camera motion if Viewer._camera_motion is not None: - Viewer._camera_xyzrpy = Viewer._camera_motion(t_simu) + Viewer._camera_xyzrpy = Viewer._camera_motion(t) # Update display if update_hook is not None: - update_hook_t = partial(update_hook, t_simu, q, v) - self.display(q, v, xyz_offset, update_hook_t, wait) + update_hook_t = partial(update_hook, t, state.q, state.v) + self.display(state.q, v, xyz_offset, update_hook_t, wait) # Sleep for a while if computing faster than display framerate sleep(1.0 / REPLAY_FRAMERATE - (time.time() - time_prev)) - - # Update time in simulation, taking into account speed ratio time_prev = time.time() - time_elapsed = time_prev - time_init - t_simu = time_interval[0] + speed_ratio * time_elapsed - # Compute corresponding right index from interpolation - i = bisect_right(times, t_simu) + # Break replay loop if the final time has been reached + if t_end - t < 1e-10: + break + + # Update simulation time, taking into account speed ratio + time_elapsed = time_prev - time_init + t = min(t_start + speed_ratio * time_elapsed, t_end) # Waiting for the first timestep is enough wait = False - - # Stop the simulation if final time is reached - if t_simu > time_interval[1]: - break except Exception as e: # Get backend info to analyze the root cause of the exception backend = Viewer.backend diff --git a/python/jiminy_py/unit_py/data/double_pendulum_hardware.toml b/python/jiminy_py/unit_py/data/double_pendulum_hardware.toml index 6d96c275c..155df5249 100644 --- a/python/jiminy_py/unit_py/data/double_pendulum_hardware.toml +++ b/python/jiminy_py/unit_py/data/double_pendulum_hardware.toml @@ -14,10 +14,10 @@ frame_name = "PendulumArm" frame_name = "SecondPendulumArm" [Sensor.EncoderSensor.PendulumJoint] -joint_name = "PendulumJoint" +motor_name = "PendulumJoint" [Sensor.EncoderSensor.SecondPendulumJoint] -joint_name = "SecondPendulumJoint" +motor_name = "SecondPendulumJoint" [Sensor.EffortSensor.PendulumJoint] motor_name = "PendulumJoint" diff --git a/python/jiminy_py/unit_py/data/foot_pendulum.toml b/python/jiminy_py/unit_py/data/foot_pendulum.toml index b991af9c6..38cdad160 100644 --- a/python/jiminy_py/unit_py/data/foot_pendulum.toml +++ b/python/jiminy_py/unit_py/data/foot_pendulum.toml @@ -36,7 +36,7 @@ frame_name = "Foot_CollisionBox_0_4" frame_name = "Foot_CollisionBox_0_6" [Sensor.EncoderSensor.PendulumJoint] -joint_name = "PendulumJoint" +motor_name = "PendulumJoint" [Sensor.EffortSensor.PendulumJoint] motor_name = "PendulumJoint" diff --git a/python/jiminy_py/unit_py/test_dense_pole.py b/python/jiminy_py/unit_py/test_dense_pole.py index 0a32d8106..e16273169 100644 --- a/python/jiminy_py/unit_py/test_dense_pole.py +++ b/python/jiminy_py/unit_py/test_dense_pole.py @@ -41,9 +41,14 @@ def setUp(self): model_options['joints']['positionLimitMin'] = [-self.joint_limit] model_options['joints']['positionLimitMax'] = [self.joint_limit] model_options['joints']['positionLimitFromUrdf'] = False - model_options['joints']['enableVelocityLimit'] = False self.robot.set_model_options(model_options) + # Configure the motors + for motor in self.robot.motors: + motor_options = motor.get_options() + motor_options["enableVelocityLimit"] = False + motor.set_options(motor_options) + # Configure the integrator engine_options = self.simulator.get_options() engine_options['stepper']['tolAbs'] = 1e-9 @@ -153,65 +158,46 @@ def test_flex_model(self): assert np.allclose(twist_flex_all[0], theta_all, atol=1e-4) def test_joint_position_limits(self): - """Test that both spring-damper and constraint models correspond to the - exact same dynamical model for joint bounds in the simple case where - the apparent inertia is constant. Then, check that joint position - limits are active when they are supposed to be. + """Test that the joint position limits are active when they are + supposed to be. """ # Define some constants t_end, step_dt = 0.05, 1e-5 - theta_all = [] - for contact_model in ('constraint', 'spring_damper'): - # Configure the engine - engine_options = self.simulator.get_options() - engine_options['stepper']['odeSolver'] = 'euler_explicit' - engine_options['stepper']['dtMax'] = step_dt - engine_options['contacts']['model'] = contact_model - self.simulator.set_options(engine_options) - - # Start the simulation - self.simulator.start(np.array((0.0,)), np.array((1.0,))) - - # Get joint bounds constraint - const = next(iter(self.robot.constraints.bounds_joints.values())) - const.kp = engine_options['joints']['boundStiffness'] - const.kd = engine_options['joints']['boundDamping'] - - # Simulate for a while - branches = set() - is_enabled = const.is_enabled - for _ in range(int(np.round(t_end / step_dt))): - self.simulator.step(step_dt) - theta = self.simulator.robot_state.q[0] - if contact_model != 'constraint': - continue - if self.joint_limit - np.abs(theta) <= 0.0: + # Configure the engine + engine_options = self.simulator.get_options() + engine_options['stepper']['odeSolver'] = 'euler_explicit' + engine_options['stepper']['dtMax'] = step_dt + self.simulator.set_options(engine_options) + + # Start the simulation + self.simulator.start(np.array((0.0,)), np.array((1.0,))) + + # Get joint bounds constraint + const = next(iter(self.robot.constraints.bounds_joints.values())) + + # Simulate for a while + branches = set() + is_enabled = const.is_enabled + for _ in range(int(np.round(t_end / step_dt))): + self.simulator.step(step_dt) + theta = self.simulator.robot_state.q[0] + if self.joint_limit - np.abs(theta) <= 0.0: + assert const.is_enabled + branches.add(0 if is_enabled else 1) + elif self.joint_limit - np.abs(theta) < self.transition_eps: + if is_enabled: assert const.is_enabled - branches.add(0 if is_enabled else 1) - elif self.joint_limit - np.abs(theta) < self.transition_eps: - if is_enabled: - assert const.is_enabled - branches.add(2) - else: - assert not const.is_enabled - branches.add(3) + branches.add(2) else: assert not const.is_enabled - branches.add(4) - is_enabled = const.is_enabled - self.simulator.stop() - if contact_model == 'constraint': - assert branches == set((0, 1, 2, 3, 4)) - - # Extract joint angle over time - log_vars = self.simulator.log_data["variables"] - (theta,) = extract_variables_from_log( - log_vars, - (f"currentPosition{self.robot.pinocchio_model.names[-1]}",)) - theta_all.append(theta) - - assert np.allclose(*theta_all, atol=1e-4) + branches.add(3) + else: + assert not const.is_enabled + branches.add(4) + is_enabled = const.is_enabled + self.simulator.stop() + assert branches == set((0, 1, 2, 3, 4)) if __name__ == '__main__': unittest.main() diff --git a/python/jiminy_py/unit_py/test_double_spring_mass.py b/python/jiminy_py/unit_py/test_double_spring_mass.py index bc34a7d1b..7cf70dec3 100644 --- a/python/jiminy_py/unit_py/test_double_spring_mass.py +++ b/python/jiminy_py/unit_py/test_double_spring_mass.py @@ -31,7 +31,7 @@ class SimulateTwoMasses(unittest.TestCase): def setUp(self): # Create the robot and load the URDF self.urdf_name = "linear_two_masses.urdf" - self.motor_names = ["FirstJoint", "SecondJoint"] + self.motor_names = ("FirstJoint", "SecondJoint") self.robot = load_urdf_default( self.urdf_name, self.motor_names, has_freeflyer=False) diff --git a/python/jiminy_py/unit_py/test_flexible_arm.py b/python/jiminy_py/unit_py/test_flexible_arm.py index 3d5b2964c..f3c42c064 100644 --- a/python/jiminy_py/unit_py/test_flexible_arm.py +++ b/python/jiminy_py/unit_py/test_flexible_arm.py @@ -95,7 +95,7 @@ def setUp(self): N_FLEXIBILITY + 1, urdf_path) # Load URDF, create model. - self.motor_names = ["base_to_link0"] + self.motor_names = ("base_to_link0",) robot = load_urdf_default(urdf_path, self.motor_names) # Remove temporary file diff --git a/python/jiminy_py/unit_py/test_foot_pendulum.py b/python/jiminy_py/unit_py/test_foot_pendulum.py index cfbdcafdd..0fc05a5c1 100644 --- a/python/jiminy_py/unit_py/test_foot_pendulum.py +++ b/python/jiminy_py/unit_py/test_foot_pendulum.py @@ -9,7 +9,7 @@ from jiminy_py.core import ( # pylint: disable=no-name-in-module - ForceSensor as force, ImuSensor as imu, ContactSensor as contact) + ForceSensor, ImuSensor, ContactSensor) from jiminy_py.simulator import Simulator @@ -58,23 +58,23 @@ def test_init_and_consistency(self): simulator.start(q0, v0) self.assertTrue(np.all(np.abs(simulator.stepper_state.a) < TOLERANCE)) - imu_data = simulator.robot.sensor_measurements[imu.type, 'Foot'] + imu_data = simulator.robot.sensor_measurements[ImuSensor.type, 'Foot'] gyro_data, accel_data = np.split(imu_data, 2) self.assertTrue(np.allclose(gyro_data, 0.0, atol=TOLERANCE)) self.assertTrue(np.allclose(accel_data, -gravity[:3], atol=TOLERANCE)) self.assertTrue(np.allclose( - simulator.robot.sensor_measurements[force.type, 'Foot'], + simulator.robot.sensor_measurements[ForceSensor.type, 'Foot'], -gravity * mass, atol=TOLERANCE)) for i in range(3): self.assertTrue(np.allclose( simulator.robot.sensor_measurements[ - contact.type, f'Foot_CollisionBox_0_{2 * i}'], + ContactSensor.type, f'Foot_CollisionBox_0_{2 * i}'], simulator.robot.sensor_measurements[ - contact.type, f'Foot_CollisionBox_0_{2 * (i + 1)}'], + ContactSensor.type, f'Foot_CollisionBox_0_{2 * (i + 1)}'], atol=TOLERANCE)) with self.assertRaises(KeyError): - simulator.robot.sensor_measurements[contact.type, 'NA'] + simulator.robot.sensor_measurements[ContactSensor.type, 'NA'] with self.assertRaises(KeyError): simulator.robot.sensor_measurements['NA', 'Foot_CollisionBox_0_1'] diff --git a/python/jiminy_py/unit_py/test_multi_robot.py b/python/jiminy_py/unit_py/test_multi_robot.py index 7735f33bc..4bdd8f267 100644 --- a/python/jiminy_py/unit_py/test_multi_robot.py +++ b/python/jiminy_py/unit_py/test_multi_robot.py @@ -27,7 +27,7 @@ def test_multi_robot(self): """ # Specify model urdf_name = "linear_single_mass.urdf" - motor_names = ["Joint"] + motor_names = ("Joint",) # Specify spring stiffness and damping for this simulation # First two parameters are the stiffness of each system, diff --git a/python/jiminy_py/unit_py/test_simple_pendulum.py b/python/jiminy_py/unit_py/test_simple_pendulum.py index d74e3b868..2b77878e7 100644 --- a/python/jiminy_py/unit_py/test_simple_pendulum.py +++ b/python/jiminy_py/unit_py/test_simple_pendulum.py @@ -3,12 +3,15 @@ integration method of jiminy on simple models. """ import unittest +from dataclasses import asdict +from typing import Union, Dict, Tuple, Sequence + import numpy as np import scipy from scipy.interpolate import interp1d -from typing import Union, Dict, Tuple, Sequence import jiminy_py.core as jiminy +from jiminy_py.log import extract_trajectory_from_log from pinocchio import Quaternion, log3, exp3 from pinocchio.rpy import matrixToRpy, rpyToMatrix @@ -30,7 +33,7 @@ class SimulateSimplePendulum(unittest.TestCase): def setUp(self): # Load URDF, create model. self.urdf_name = "simple_pendulum.urdf" - self.motor_names = ["PendulumJoint"] + self.motor_names = ("PendulumJoint",) self.robot = load_urdf_default(self.urdf_name, self.motor_names) # Add IMU to the robot @@ -137,6 +140,103 @@ def spring_force(t, q, v, sensor_measurements, u_custom): self.assertTrue(np.allclose(x_jiminy, x_analytical, atol=TOLERANCE)) + def test_velocity_bounds(self): + VELOCITY_MAX = 15.0 + VELOCITY_SLOPE = 0.5 + TAU = 50.0 + + # Configure the robot: set rotor inertia + robot_options = self.robot.get_options() + motor_options = robot_options["motors"]["PendulumJoint"] + motor_options['enableEffortLimit'] = True + motor_options['enableVelocityLimit'] = True + motor_options['velocityEffortInvSlope'] = VELOCITY_SLOPE + motor_options['velocityLimitFromUrdf'] = False + motor_options['velocityLimit'] = VELOCITY_MAX + self.robot.set_options(robot_options) + + # Initialize the controller and setup the engine + def compute_command(t, q, v, sensor_measurements, command): + command[:] = TAU + + engine = jiminy.Engine() + setup_controller_and_engine( + engine, self.robot, compute_command=compute_command) + + # Disable position bounds + self.robot.pinocchio_model_th.lowerPositionLimit[:] = -1000.0 + self.robot.pinocchio_model_th.upperPositionLimit[:] = +1000.0 + + # Configure the engine + engine_options = engine.get_options() + engine_options["world"]["gravity"] = np.zeros(6) + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["tolAbs"] = 1e-12 + engine_options["stepper"]["tolRel"] = 1e-12 + engine.set_options(engine_options) + + # Run simulation + engine.simulate(4.0, np.array([0.0]), np.array([0.0])) + + # Extract log data + log_data = engine.log_data + time = engine.log_data['variables']['Global.Time'] + vel = engine.log_data['variables']['currentVelocityPendulum'] + acc = engine.log_data['variables']['currentAccelerationPendulum'] + + # Check that the velocity does not exceed the limit + assert np.all(np.abs(vel) < VELOCITY_MAX) + + # Check that the velocity ultimately reaches the limit + assert np.abs(vel[-1]) - VELOCITY_MAX < TOLERANCE + + # Check that the acceleration is decreasing exponentially at some point + acc_thr = TAU / self.robot.pinocchio_data.mass[0] + start_idx = next(iter(i for i, a in enumerate(acc) if a < acc_thr)) + end_idx = start_idx + next(iter(i for i, a in enumerate(acc) if a < 0.1)) + acc_slope = np.diff( + np.log(acc[start_idx:end_idx] / acc_thr) + ) / np.diff(time[start_idx:end_idx]) + assert np.all(np.abs(acc_slope - np.mean(acc_slope)) < 1e-5) + + # Check that the velocity limit activates when it is expected + motor = self.robot.motors[0] + v_th_max = max( + motor.velocity_limit - VELOCITY_SLOPE * motor.effort_limit, 0.0) + v_th = motor.velocity_limit - ( + TAU / motor.effort_limit) * (motor.velocity_limit - v_th_max) + assert vel[start_idx - 1] < v_th and vel[start_idx] > v_th + + # Check that the acceleration cancels out completely + assert np.abs(acc[-1]) < TOLERANCE + + def test_extract_trajectory(self): + """TODO: Write documentation. + """ + # Instantiate the simulator + engine = jiminy.Engine() + setup_controller_and_engine(engine, self.robot) + robot_state = engine.robot_states[0] + + # Run a simulation + T_END = 4.0 + engine.simulate(T_END, np.array([0.0]), np.array([1.0])) + + # Extract trajectory from log + traj = extract_trajectory_from_log(engine.log_data) + + # Check that the trajectory is consistent with the final state + state_0, *_, state_f = traj.states + assert np.abs(state_0.t) < 1e-12 + assert np.abs(state_f.t - T_END) < 1e-12 + for key, value in asdict(state_f).items(): + if key == "t": + continue + data = getattr(robot_state, key) + if key == "f_external": + data = tuple(f_ext.vector for f_ext in data) + np.testing.assert_allclose(data, value, rtol=0, atol=1e-12) + def test_pendulum_integration(self): """Compare pendulum motion, as simulated by Jiminy, against an equivalent simulation done in python. @@ -184,12 +284,12 @@ def test_backlash(self): self.robot.set_options(robot_options) TAU = 5.0 - def ControllerConstant(t, q, v, sensor_measurements, command): + def compute_command(t, q, v, sensor_measurements, command): command[:] = - TAU engine = jiminy.Engine() setup_controller_and_engine( - engine, self.robot, compute_command=ControllerConstant) + engine, self.robot, compute_command=compute_command) engine_options = engine.get_options() engine_options["constraints"]["regularization"] = 0.0 diff --git a/python/jiminy_py/unit_py/test_simulator.py b/python/jiminy_py/unit_py/test_simulator.py index 75486f538..4fc3921e6 100644 --- a/python/jiminy_py/unit_py/test_simulator.py +++ b/python/jiminy_py/unit_py/test_simulator.py @@ -1,4 +1,4 @@ -"""TODO: Write documentation. +""" TODO: Write documentation. """ import os import tempfile @@ -7,6 +7,8 @@ import numpy as np import jiminy_py.core as jiminy +from jiminy_py.core import ImuSensor # pylint: disable=no-name-in-module + from jiminy_py.robot import BaseJiminyRobot from jiminy_py.simulator import Simulator @@ -16,6 +18,8 @@ extract_trajectories_from_log) from jiminy_py.viewer.replay import play_logs_data +from utilities import setup_controller_and_engine + class SimulatorTest(unittest.TestCase): def test_consistency_velocity_acceleration(self): @@ -48,16 +52,13 @@ def test_consistency_velocity_acceleration(self): imu.initialize(fname) # Define a PD controller with fixed target position - class Controller(jiminy.BaseController): - def compute_command(self, t, q, v, command): - target = np.array([1.5, 0.0]) - command[:] = -5000 * ((q[::2] - target) + 0.07 * v[::2]) - - robot.controller = Controller() + def compute_command(t, q, v, sensor_measurements, command): + target = np.array([1.5, 0.0]) + command[:] = -5000 * ((q[::2] - target) + 0.07 * v[::2]) # Instantiate the engine engine = jiminy.Engine() - engine.add_robot(robot) + setup_controller_and_engine(engine, robot, compute_command) # Configuration the simulation engine_options = engine.get_options() @@ -95,11 +96,10 @@ def compute_command(self, t, q, v, command): diff_velocities, accelerations[:, :-1], atol=1e-12, rtol=0.0) # Check that IMU accelerations match gravity at rest - for imu_name in robot.sensor_names[jiminy.ImuSensor.type]: - log_imu_name = ".".join((jiminy.ImuSensor.type, imu_name)) + for imu_sensor in robot.sensors[ImuSensor.type]: + log_imu_name = ".".join((ImuSensor.type, imu_sensor.name)) imu_data = np.stack(extract_variables_from_log( - log_vars, jiminy.ImuSensor.fieldnames, log_imu_name - ), axis=0) + log_vars, ImuSensor.fieldnames, log_imu_name), axis=0) imu_gyro, imu_accel = np.split(imu_data, 2) imu_gyro_norm = np.linalg.norm(imu_gyro, axis=0) imu_accel_norm = np.linalg.norm(imu_accel, axis=0) @@ -157,7 +157,7 @@ def test_single_robot_simulation(self): log_data = read_log(log_path) trajectory = extract_trajectory_from_log(log_data) - final_log_states = trajectory['evolution_robot'][-1] + final_log_states = trajectory.states[-1] np.testing.assert_array_almost_equal( simulator.robot_state.q, final_log_states.q, decimal=10) @@ -169,7 +169,7 @@ def test_single_robot_simulation(self): tempfile.gettempdir(), f"video_{next(tempfile._get_candidate_names())}.mp4") play_logs_data( - robot, log_data, record_video_path=video_path, verbose=False) + log_data, robot, record_video_path=video_path, verbose=False) self.assertTrue(os.path.isfile(video_path)) def test_double_robot_simulation(self): @@ -225,7 +225,7 @@ def test_double_robot_simulation(self): trajectories = extract_trajectories_from_log(log_data) trajectory_1, trajectory_2 = ( - trajectories[robot.name]['evolution_robot'][-1] + trajectories[robot.name].states[-1] for robot in simulator.robots) robot_states_1, robot_states_2 = simulator.robot_states diff --git a/python/jiminy_py/unit_py/utilities.py b/python/jiminy_py/unit_py/utilities.py index dc557f510..28dd8055b 100644 --- a/python/jiminy_py/unit_py/utilities.py +++ b/python/jiminy_py/unit_py/utilities.py @@ -49,10 +49,11 @@ def load_urdf_default(urdf_name: str, # Configure the robot robot_options = robot.get_options() - robot_options["model"]["joints"]["enableVelocityLimit"] = False - for name in robot.motor_names: - robot_options["motors"][name]['enableCommandLimit'] = False - robot_options["motors"][name]['enableArmature'] = False + motor_options = robot_options["motors"] + for motor in robot.motors: + motor_options[motor.name]["enableVelocityLimit"] = False + motor_options[motor.name]['enableEffortLimit'] = False + motor_options[motor.name]['enableArmature'] = False robot.set_options(robot_options) return robot @@ -112,6 +113,15 @@ def setup_controller_and_engine( # Initialize the engine engine.add_robot(robot) + # Enable telemetry by default + engine_options = engine.get_options() + telemetry_options = engine_options["telemetry"] + telemetry_options["enableEffort"] = True + telemetry_options["enableCommand"] = True + telemetry_options["enableForceExternal"] = True + telemetry_options["enableEnergy"] = True + engine.set_options(engine_options) + def neutral_state(robot: jiminy.Model, split: bool = False @@ -188,6 +198,15 @@ def simulate_and_get_state_evolution( associated sequence of states as a 2D array each line corresponds to a given time. """ + # Enable telemetry by default + engine_options = engine.get_options() + telemetry_options = engine_options["telemetry"] + telemetry_options["enableEffort"] = True + telemetry_options["enableCommand"] = True + telemetry_options["enableForceExternal"] = True + telemetry_options["enableEnergy"] = True + engine.set_options(engine_options) + # Run simulation if isinstance(x0, np.ndarray): robot, = engine.robots diff --git a/python/jiminy_pywrap/include/jiminy/python/utilities.h b/python/jiminy_pywrap/include/jiminy/python/utilities.h index bd367bb42..de6b4f812 100644 --- a/python/jiminy_pywrap/include/jiminy/python/utilities.h +++ b/python/jiminy_pywrap/include/jiminy/python/utilities.h @@ -19,14 +19,17 @@ #include -#define EXPECTED_PYTYPE_FOR_ARG_IS_ARRAY(type) \ - namespace boost::python::converter \ - { \ - template<> \ - struct expected_pytype_for_arg \ - { \ - static const PyTypeObject * get_pytype() { return &PyArray_Type; } \ - }; \ +#define EXPECTED_PYTYPE_FOR_ARG_IS_ARRAY(type) \ + namespace boost::python::converter \ + { \ + template<> \ + struct expected_pytype_for_arg \ + { \ + static const PyTypeObject * get_pytype() \ + { \ + return &PyArray_Type; \ + } \ + }; \ } EXPECTED_PYTYPE_FOR_ARG_IS_ARRAY(numpy::ndarray) @@ -217,6 +220,7 @@ namespace jiminy::python { stringStream << ":\n\n" << doc; } + // FIXME: Use `view` to get a string_view rather than a string copy when moving to C++20 return stringStream.str().substr( std::min(static_cast(stringStream.tellp()), size_t(1))); } @@ -235,7 +239,7 @@ namespace jiminy::python doc, std::pair{"fget", getMemberFuncPtr}, std::pair{"fset", setMemberFuncPtr}); } -// clang-format off + // clang-format off #define DEF_READONLY3(namePy, memberFuncPtr, doc) \ def_readonly(namePy, \ memberFuncPtr, \ @@ -628,11 +632,23 @@ namespace jiminy::python "'values' input array must have dtype 'np.float64' or 'np.int64'."); } + template + struct is_boost_variant : std::false_type + { + }; + + template + struct is_boost_variant> : std::true_type + { + }; + + template + inline constexpr bool is_boost_variant_v = is_boost_variant...>::value; + /// Convert most C++ objects into Python objects by value template - std::enable_if_t && !is_array_v && !is_eigen_any_v && - !std::is_arithmetic_v> && - !std::is_same_v, GenericConfig> && + std::enable_if_t && !is_array_v && !is_eigen_any_v && !is_map_v && + !is_boost_variant_v && !std::is_arithmetic_v> && !std::is_same_v, SensorMeasurementTree::value_type> && !std::is_same_v, std::string_view> && !std::is_same_v, FlexibilityJointConfig>, @@ -693,7 +709,7 @@ namespace jiminy::python template std::enable_if_t, FlexibilityJointConfig>, bp::object> - convertToPython(T && flexibilityJointConfig, const bool & copy) + convertToPython(T && flexibilityJointConfig, const bool & copy = true) { if (!copy) { @@ -712,28 +728,40 @@ namespace jiminy::python template std::enable_if_t, SensorMeasurementTree::value_type>, bp::object> - convertToPython(T && sensorDataTypeItem, const bool & copy) + convertToPython(T && sensorDataTypeItem, const bool & copy = true) { auto & [sensorGroupName, sensorDataType] = sensorDataTypeItem; return bp::make_tuple(sensorGroupName, convertToPython(sensorDataType.getAll(), copy)); } + template + std::enable_if_t, bp::object> convertToPython(T && value, + const bool & copy = true); + template std::enable_if_t || is_array_v, bp::object> - convertToPython(T && data, const bool & copy = true) + convertToPython(T && vect, const bool & copy = true) { - bp::list dataPy; - for (auto & val : data) + bp::list listPy; + for (auto && value : vect) { - dataPy.append(convertToPython(val, copy)); + listPy.append(convertToPython(value, copy)); } // FIXME: Remove explicit and redundant move when moving to C++20 - return std::move(dataPy); + return std::move(listPy); } template - std::enable_if_t, GenericConfig>, bp::object> - convertToPython(T && config, const bool & copy = true); + std::enable_if_t, bp::object> convertToPython(T && dict, const bool & copy = true) + { + bp::dict dictPy; + for (auto && [key, value] : dict) + { + dictPy[key] = convertToPython(value, copy); + } + // FIXME: Remove explicit and redundant move when moving to C++20 + return std::move(dictPy); + } class AppendBoostVariantToPython : public boost::static_visitor { @@ -754,17 +782,10 @@ namespace jiminy::python }; template - std::enable_if_t, GenericConfig>, bp::object> - convertToPython(T && config, const bool & copy) + std::enable_if_t, bp::object> convertToPython(T && value, + const bool & copy) { - bp::dict configPyDict; - AppendBoostVariantToPython visitor(copy); - for (auto & [key, value] : config) - { - configPyDict[key] = boost::apply_visitor(visitor, value); - } - // FIXME: Remove explicit and redundant move when moving to C++20 - return std::move(configPyDict); + return boost::apply_visitor(AppendBoostVariantToPython{copy}, value); } template @@ -778,7 +799,7 @@ namespace jiminy::python { return &PyList_Type; } - else if constexpr (std::is_same_v, GenericConfig> || + else if constexpr (is_map_v || std::is_same_v, FlexibilityJointConfig>) { return &PyDict_Type; @@ -1158,8 +1179,8 @@ namespace jiminy::python }; template - ConvertGeneratorFromPythonAndInvoke(R (*)(Args...)) - -> ConvertGeneratorFromPythonAndInvoke; + ConvertGeneratorFromPythonAndInvoke( + R (*)(Args...)) -> ConvertGeneratorFromPythonAndInvoke; template class ConvertGeneratorFromPythonAndInvoke @@ -1185,8 +1206,8 @@ namespace jiminy::python }; template - ConvertGeneratorFromPythonAndInvoke(R (T::*)(Args...)) - -> ConvertGeneratorFromPythonAndInvoke; + ConvertGeneratorFromPythonAndInvoke( + R (T::*)(Args...)) -> ConvertGeneratorFromPythonAndInvoke; // **************************** Automatic From Python converter **************************** // diff --git a/python/jiminy_pywrap/src/engine.cc b/python/jiminy_pywrap/src/engine.cc index cb7dadfb4..ee6e2a148 100644 --- a/python/jiminy_pywrap/src/engine.cc +++ b/python/jiminy_pywrap/src/engine.cc @@ -156,6 +156,7 @@ namespace jiminy::python s << "command:\n " << self.command.transpose().format(HeavyFmt); s << "u:\n " << self.u.transpose().format(HeavyFmt); s << "u_motor:\n " << self.uMotor.transpose().format(HeavyFmt); + s << "u_transmission:\n " << self.uTransmission.transpose().format(HeavyFmt); s << "u_internal:\n " << self.uInternal.transpose().format(HeavyFmt); s << "u_custom:\n " << self.uCustom.transpose().format(HeavyFmt); s << "f_external:\n"; @@ -178,6 +179,7 @@ namespace jiminy::python .DEF_READONLY("command", &RobotState::command) .DEF_READONLY("u", &RobotState::u) .DEF_READONLY("u_motor", &RobotState::uMotor) + .DEF_READONLY("u_transmission", &RobotState::uTransmission) .DEF_READONLY("u_internal", &RobotState::uInternal) .DEF_READONLY("u_custom", &RobotState::uCustom) .DEF_READONLY("f_external", &RobotState::fExternal) diff --git a/python/jiminy_pywrap/src/generators.cc b/python/jiminy_pywrap/src/generators.cc index 50fd00b12..725846a55 100644 --- a/python/jiminy_pywrap/src/generators.cc +++ b/python/jiminy_pywrap/src/generators.cc @@ -193,6 +193,7 @@ namespace jiminy::python bp::def("random_tile_ground", &tiles, (bp::arg("size"), "height_max", "interp_delta", "sparsity", "orientation", "seed")); + bp::def("stairs_ground", &stairs, (bp::arg("step_width"), "step_height", "step_number", "orientation")); bp::def("sum_heightmaps", &sumHeightmaps, (bp::arg("heightmaps"))); bp::def("merge_heightmaps", &mergeHeightmaps, (bp::arg("heightmaps"))); diff --git a/python/jiminy_pywrap/src/helpers.cc b/python/jiminy_pywrap/src/helpers.cc index 73256b94d..2789db25a 100644 --- a/python/jiminy_pywrap/src/helpers.cc +++ b/python/jiminy_pywrap/src/helpers.cc @@ -68,6 +68,12 @@ namespace jiminy::python return bp::make_tuple(model, collisionModel); } + bp::object saveRobotToBinary(const std::shared_ptr & robot) + { + std::string data = saveToBinary(robot, true); + return bp::object(bp::handle<>(PyBytes_FromStringAndSize(data.c_str(), data.size()))); + } + std::shared_ptr buildRobotFromBinary(const std::string & data, const bp::object & meshPathDirPy, const bp::object & packageDirsPy) @@ -328,11 +334,14 @@ namespace jiminy::python bp::arg("build_visual_model") = false, bp::arg("load_visual_meshes") = false)); - bp::def("build_robot_from_binary", &buildRobotFromBinary, - (bp::arg("data"), - bp::arg("mesh_path_dir") = bp::object(), - bp::arg("mesh_package_dirs") = bp::list())); + bp::def("save_to_binary", &saveRobotToBinary, (bp::arg("robot"))); + + bp::def("load_from_binary", &buildRobotFromBinary, + (bp::arg("data"), + bp::arg("mesh_path_dir") = bp::object(), + bp::arg("mesh_package_dirs") = bp::list())); + bp::def("get_joint_type", &getJointType, bp::arg("joint_model")); bp::def("get_joint_type", &getJointTypeFromIndex, (bp::arg("pinocchio_model"), "joint_index")); bp::def("get_joint_indices", &getJointIndices, diff --git a/python/jiminy_pywrap/src/motors.cc b/python/jiminy_pywrap/src/motors.cc index 1cfaf9ab6..c7c1b94a0 100644 --- a/python/jiminy_pywrap/src/motors.cc +++ b/python/jiminy_pywrap/src/motors.cc @@ -36,10 +36,10 @@ namespace jiminy::python &AbstractMotorBase::getJointName, bp::return_value_policy()) .ADD_PROPERTY_GET("joint_index", &AbstractMotorBase::getJointIndex) - .ADD_PROPERTY_GET("joint_type", &AbstractMotorBase::getJointType) - .ADD_PROPERTY_GET("joint_position_index", &AbstractMotorBase::getJointPositionIndex) - .ADD_PROPERTY_GET("joint_velocity_index", &AbstractMotorBase::getJointVelocityIndex) - .ADD_PROPERTY_GET("command_limit", &AbstractMotorBase::getCommandLimit) + .ADD_PROPERTY_GET("position_limit_lower", &AbstractMotorBase::getPositionLimitLower) + .ADD_PROPERTY_GET("position_limit_upper", &AbstractMotorBase::getPositionLimitUpper) + .ADD_PROPERTY_GET("velocity_limit", &AbstractMotorBase::getVelocityLimit) + .ADD_PROPERTY_GET("effort_limit", &AbstractMotorBase::getEffortLimit) .ADD_PROPERTY_GET("armature", &AbstractMotorBase::getArmature) .ADD_PROPERTY_GET("backlash", &AbstractMotorBase::getBacklash) diff --git a/python/jiminy_pywrap/src/robot.cc b/python/jiminy_pywrap/src/robot.cc index da2577e46..34a2a8798 100644 --- a/python/jiminy_pywrap/src/robot.cc +++ b/python/jiminy_pywrap/src/robot.cc @@ -7,6 +7,7 @@ #include "jiminy/core/robot/robot.h" #include "pinocchio/bindings/python/fwd.hpp" +#include "pinocchio/bindings/python/utils/copyable.hpp" #include "jiminy/python/utilities.h" #include "jiminy/python/robot.h" @@ -166,6 +167,8 @@ namespace jiminy::python void exposeModel() { bp::class_, boost::noncopyable>("Model", bp::init<>()) + .def("copy", &std::make_shared, bp::arg("self")) + .def("initialize", &initialize, (bp::arg("self"), @@ -316,16 +319,6 @@ namespace jiminy::python &Model::getBacklashJointIndices, bp::return_value_policy>()) - .ADD_PROPERTY_GET_WITH_POLICY("position_limit_lower", - &Model::getPositionLimitMin, - bp::return_value_policy()) - .ADD_PROPERTY_GET_WITH_POLICY("position_limit_upper", - &Model::getPositionLimitMax, - bp::return_value_policy()) - .ADD_PROPERTY_GET_WITH_POLICY("velocity_limit", - &Model::getVelocityLimit, - bp::return_value_policy()) - .ADD_PROPERTY_GET_WITH_POLICY("log_position_fieldnames", &Model::getLogPositionFieldnames, bp::return_value_policy>()) @@ -335,6 +328,9 @@ namespace jiminy::python .ADD_PROPERTY_GET_WITH_POLICY("log_acceleration_fieldnames", &Model::getLogAccelerationFieldnames, bp::return_value_policy>()) + .ADD_PROPERTY_GET_WITH_POLICY("log_effort_fieldnames", + &Model::getLogEffortFieldnames, + bp::return_value_policy>()) .ADD_PROPERTY_GET_WITH_POLICY("log_f_external_fieldnames", &Model::getLogForceExternalFieldnames, bp::return_value_policy>()); @@ -355,17 +351,6 @@ namespace jiminy::python return std::make_shared(self.getSensorMeasurements()); } - static bp::dict getSensorNames(Robot & self) - { - bp::dict sensorsNamesPy; - const auto & sensorsNames = self.getSensorNames(); - for (const auto & sensorTypeNames : sensorsNames) - { - sensorsNamesPy[sensorTypeNames.first] = convertToPython(sensorTypeNames.second); - } - return sensorsNamesPy; - } - static void setModelOptions(Robot & self, const bp::dict & configPy) { GenericConfig config = self.getModelOptions(); @@ -378,6 +363,8 @@ namespace jiminy::python { bp::class_, std::shared_ptr, boost::noncopyable>( "Robot", bp::init(bp::arg("name") = "")) + .def("copy", &std::make_shared, bp::arg("self")) + .def("initialize", &initialize, (bp::arg("self"), @@ -398,10 +385,6 @@ namespace jiminy::python "name", &Robot::getName, bp::return_value_policy()) .def("attach_motor", &Robot::attachMotor, (bp::arg("self"), "motor")) - .def("get_motor", - static_cast (Robot::*)(const std::string &)>( - &Robot::getMotor), - (bp::arg("self"), "motor_name")) .def("detach_motor", static_cast(&Robot::detachMotor), (bp::arg("self"), "joint_name")) @@ -415,10 +398,23 @@ namespace jiminy::python .def("detach_sensors", &Robot::detachSensors, (bp::arg("self"), bp::arg("sensor_type") = std::string())) + + .def("get_motor", + static_cast (Robot::*)(const std::string &)>( + &Robot::getMotor), + (bp::arg("self"), "motor_name")) + .ADD_PROPERTY_GET_WITH_POLICY( + "motors", + static_cast(&Robot::getMotors), + bp::return_value_policy>()) .def("get_sensor", static_cast (Robot::*)( const std::string &, const std::string &)>(&Robot::getSensor), (bp::arg("self"), "sensor_type", "sensor_name")) + .ADD_PROPERTY_GET_WITH_POLICY( + "sensors", + static_cast(&Robot::getSensors), + bp::return_value_policy>()) .ADD_PROPERTY_GET_SET("controller", static_cast (Robot::*)()>( @@ -439,27 +435,10 @@ namespace jiminy::python &Robot::getModelOptions, bp::return_value_policy()) - .ADD_PROPERTY_GET("nmotors", &Robot::nmotors) - .ADD_PROPERTY_GET_WITH_POLICY("motor_names", - &Robot::getMotorNames, - bp::return_value_policy>()) - .ADD_PROPERTY_GET_WITH_POLICY("motor_position_indices", - &Robot::getMotorsPositionIndices, - bp::return_value_policy>()) - .ADD_PROPERTY_GET_WITH_POLICY("motor_velocity_indices", - &Robot::getMotorVelocityIndices, - bp::return_value_policy>()) - .ADD_PROPERTY_GET("sensor_names", &internal::robot::getSensorNames) - - .ADD_PROPERTY_GET_WITH_POLICY("command_limit", - &Robot::getCommandLimit, - bp::return_value_policy()) - .ADD_PROPERTY_GET_WITH_POLICY("log_command_fieldnames", &Robot::getLogCommandFieldnames, bp::return_value_policy>()) - .ADD_PROPERTY_GET_WITH_POLICY("log_motor_effort_fieldnames", - &Robot::getLogMotorEffortFieldnames, - bp::return_value_policy>()); + + .ADD_PROPERTY_GET("nmotors", &Robot::nmotors); } } diff --git a/python/jiminy_pywrap/src/sensors.cc b/python/jiminy_pywrap/src/sensors.cc index e3bb51356..234785b0b 100644 --- a/python/jiminy_pywrap/src/sensors.cc +++ b/python/jiminy_pywrap/src/sensors.cc @@ -279,6 +279,30 @@ namespace jiminy::python // ************************************** BasicSensors ************************************* // + namespace internal::encoder + { + void Initialize( + EncoderSensor & self, const bp::object & motorNamePy, const bp::object & jointNamePy) + { + if (!(motorNamePy.is_none() ^ jointNamePy.is_none())) + { + throw std::invalid_argument( + "Either 'motor_name' or 'joint_name' must be specified but not both."); + } + + if (jointNamePy.is_none()) + { + const std::string motorName = bp::extract(motorNamePy); + self.initialize(motorName, false); + } + else + { + const std::string jointName = bp::extract(jointNamePy); + self.initialize(jointName, true); + } + } + } + struct PyBasicSensorsVisitor : public bp::def_visitor { public: @@ -350,12 +374,19 @@ namespace jiminy::python boost::noncopyable>( "EncoderSensor", bp::init((bp::arg("self"), "name"))) .def(PyBasicSensorsVisitor()) - .def("initialize", &EncoderSensor::initialize, (bp::arg("self"), "joint_name")) + .def("initialize", + &internal::encoder::Initialize, + (bp::arg("self"), + bp::arg("motor_name") = bp::object(), + bp::arg("joint_name") = bp::object())) .ADD_PROPERTY_GET_WITH_POLICY("joint_name", &EncoderSensor::getJointName, bp::return_value_policy()) .ADD_PROPERTY_GET("joint_index", &EncoderSensor::getJointIndex) - .ADD_PROPERTY_GET("joint_type", &EncoderSensor::getJointType); + .ADD_PROPERTY_GET_WITH_POLICY("motor_name", + &EncoderSensor::getMotorName, + bp::return_value_policy()) + .ADD_PROPERTY_GET("motor_index", &EncoderSensor::getMotorIndex); bp::class_,