diff --git a/README.md b/README.md index 47156100..9bd70c7f 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ **Robot Control Stack (RCS)** is a flexible, native [Gymnasium](https://gymnasium.farama.org/) wrapper-based robot control interface designed specifically for modern robot learning and Vision-Language-Action (VLA) models. -It completely unifies **MuJoCo simulation** and real-world physical robot control into a single, seamless API. Currently, RCS natively supports four robots out-of-the-box: **Franka FR3/Panda, xArm7, UR5e, and SO101.** +It completely unifies **MuJoCo simulation** and real-world physical robot control into a single, seamless API. Currently, RCS natively supports five robots out-of-the-box: **Franka FR3/Panda, xArm7, UR5e, SO101, and I2RT YAM.** ![RCS Demo](https://raw.githubusercontent.com/RobotControlStack/robotcontrolstack.github.io/refs/heads/master/static/videos/grid.webp) @@ -26,7 +26,7 @@ Traditional robotics middleware (like ROS/ROS2) and complex motion planning pipe * **Zero ROS Overhead:** No complex message-passing, middleware, or network configuration required. Run natively in Python with a lightweight C++ backend. * **Frictionless Sim-to-Real:** Train your Reinforcement Learning or VLA policies in our MuJoCo Gymnasium wrapper, and deploy the *exact same code* directly to physical hardware. * **Synchronous Execution:** Optimized specifically for the highly parallelized, synchronous data collection required by modern ML workflows. -* **Ready-to-Use Apps:** Ships with pre-built applications for data collection via teleoperation and remote model inference via [vlagents](https://github.com/RobotControlStack/vlagents). See [examples/teleop/README.md](examples/teleop/README.md) and [examples/inference/README.md](examples/inference/README.md). +* **Ready-to-Use Apps:** Ships with pre-built applications for data collection via teleoperation and remote model inference via [vlagents](https://github.com/RobotControlStack/vlagents). See the [teleoperation guide](examples/teleop/README.md), and [inference guide](examples/inference/README.md). ## 🧩 Wrapper-Based Architecture @@ -168,7 +168,7 @@ export RCS_PREFIX=/path/to/rcs-assets ## 🦾 Hardware Extensions -RCS supports various hardware extensions to seamlessly connect your policies to the real world (e.g., FR3, xArm7, RealSense). These are located in the `extensions` directory. +RCS supports various hardware extensions to seamlessly connect your policies to the real world (e.g., FR3, xArm7, YAM, RealSense). These are located in the `extensions` directory. To install a specific robot extension (example for Franka FR3): diff --git a/docs/apps/index.md b/docs/apps/index.md index 799132db..a1cb72bc 100644 --- a/docs/apps/index.md +++ b/docs/apps/index.md @@ -4,12 +4,13 @@ RCS ships with ready-to-use applications for common operator workflows such as r ## Teleoperation -Use the Franka teleoperation app when you want to collect demonstrations or directly control a robot from an operator interface. +Use the Meta Quest 3 teleoperation examples when you want to collect demonstrations or directly control a robot from an operator interface. RCS provides examples for Franka and I2RT YAM arms. - Example guide: [examples/teleop/README.md](../../examples/teleop/README.md) -- Main script: [examples/teleop/franka.py](../../examples/teleop/franka.py) +- Franka example: [examples/teleop/franka.py](../../examples/teleop/franka.py) +- YAM example: [examples/teleop/yam.py](../../examples/teleop/yam.py) -The current example focuses on Franka teleoperation with Meta Quest 3 and GELLO-based setups. +The Franka example supports Meta Quest 3 and GELLO-based setups. The YAM example supports dual-arm Meta Quest 3 teleoperation in simulation or on hardware; install the [YAM extension](../extensions/rcs_yam.md) before using real YAM hardware. ## Inference diff --git a/docs/extensions/rcs_yam.md b/docs/extensions/rcs_yam.md index d8599e12..dee5b74e 100644 --- a/docs/extensions/rcs_yam.md +++ b/docs/extensions/rcs_yam.md @@ -58,3 +58,9 @@ See `extensions/rcs_yam/README.md` for the full extension documentation and `extensions/rcs_yam/src/rcs_yam/scripts/test_robot.py` for a bring-up script. For a maintained example, see `examples/yam/yam_env_cartesian_control.py`, which moves the TCP forward and backward in synchronous Cartesian mode in simulation or on hardware. + +## Teleoperation example + +Use [examples/teleop/yam.py](../../examples/teleop/yam.py) to teleoperate YAM arms with a Meta Quest +3. The example runs against simulation or hardware, and can optionally record RealSense cameras. See +the [teleoperation README](../../examples/teleop/README.md) for setup instructions. diff --git a/extensions/rcs_fr3/src/hw/Franka.cpp b/extensions/rcs_fr3/src/hw/Franka.cpp index 4ea7ccda..efdc9d11 100644 --- a/extensions/rcs_fr3/src/hw/Franka.cpp +++ b/extensions/rcs_fr3/src/hw/Franka.cpp @@ -19,6 +19,16 @@ namespace rcs { namespace hw { +common::Pose GetFlangeInBaseFrame(const franka::RobotState& robot_state) { + return common::Pose(robot_state.O_T_EE) * + common::Pose(robot_state.F_T_EE).inverse(); +} + +common::Pose GetTCPInBaseFrame(const franka::RobotState& robot_state, + const common::Pose& tcp_offset) { + return GetFlangeInBaseFrame(robot_state) * tcp_offset; +} + Franka::Franka(const FrankaConfig& cfg, std::optional> ik) : m_cfg(cfg), @@ -99,19 +109,29 @@ void Franka::set_default_robot_behavior() { common::Pose Franka::get_cartesian_position() { this->check_for_background_errors(); - common::Pose x; + franka::RobotState robot_state; if (this->running_controller.load() == Controller::none) { this->curr_state = this->robot.readOnce(); - x = common::Pose(this->curr_state.O_T_EE); + robot_state = this->curr_state; } else { this->interpolator_mutex.lock(); - x = common::Pose(this->curr_state.O_T_EE); + robot_state = this->curr_state; this->interpolator_mutex.unlock(); } - if (!this->m_cfg.tcp_offset_configured_in_desk) { - return x * this->m_cfg.tcp_offset; + return GetTCPInBaseFrame(robot_state, this->m_cfg.tcp_offset); +} + +common::Pose Franka::get_cartesian_flange_position() { + this->check_for_background_errors(); + franka::RobotState robot_state; + if (this->running_controller.load() == Controller::none) { + this->curr_state = this->robot.readOnce(); + robot_state = this->curr_state; + } else { + std::lock_guard lock(this->interpolator_mutex); + robot_state = this->curr_state; } - return x; + return GetFlangeInBaseFrame(robot_state); } void Franka::set_joint_position(const common::VectorXd& q) { @@ -164,7 +184,7 @@ void PInverse(const Eigen::MatrixXd& M, Eigen::MatrixXd& M_inv, } void TorqueSafetyGuardFn(std::array& tau_d_array, - const std::array& torque_limit) { + const common::Vector7d& torque_limit) { for (size_t i = 0; i < tau_d_array.size(); i++) { if (tau_d_array[i] < -torque_limit[i]) { tau_d_array[i] = -torque_limit[i]; @@ -250,10 +270,8 @@ void Franka::osc_set_cartesian_position( this->interpolator_mutex.lock(); } - common::Pose curr_pose(this->curr_state.O_T_EE); - if (!this->m_cfg.tcp_offset_configured_in_desk) { - curr_pose = curr_pose * this->m_cfg.tcp_offset; - } + common::Pose curr_pose = + GetTCPInBaseFrame(this->curr_state, this->m_cfg.tcp_offset); this->traj_interpolator.reset( this->controller_time, curr_pose.translation(), curr_pose.quaternion(), desired_pose_EE_in_base_frame.translation(), @@ -284,18 +302,22 @@ void Franka::osc() { franka::Model model = this->robot.loadModel(); const Eigen::Vector3d kp_p_cfg = this->m_cfg.kp_p; const double kp_r_cfg = this->m_cfg.kp_r; + const common::Vector7d torque_limit = this->m_cfg.torque_limit; + const bool allow_high_collision = this->m_cfg.allow_high_collision; this->controller_time = 0.0; // conservative collision and impedance behavior this->set_default_robot_behavior(); - // high collision threshold values for high impedance - this->robot.setCollisionBehavior( - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + if (allow_high_collision) { + // High collision threshold values for high impedance. + this->robot.setCollisionBehavior( + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + } // from bench mark // ([150.0, 150.0, 60.0], 250.0), // kp_translation, kp_rotation @@ -381,8 +403,9 @@ void Franka::osc() { Eigen::Map> gravity( gravity_array.data()); - std::array jacobian_array = - model.zeroJacobian(franka::Frame::kEndEffector, robot_state); + std::array jacobian_array = model.zeroJacobian( + franka::Frame::kEndEffector, robot_state.q, + this->m_cfg.tcp_offset.affine_array(), robot_state.EE_T_K); Eigen::Map> jacobian( jacobian_array.data()); @@ -394,9 +417,7 @@ void Franka::osc() { // Express OSC feedback in the same TCP frame exposed by the public // Cartesian API. common::Pose T_EE_in_base_frame_pose = - this->m_cfg.tcp_offset_configured_in_desk - ? common::Pose(robot_state.O_T_EE) - : common::Pose(robot_state.O_T_EE) * this->m_cfg.tcp_offset; + GetTCPInBaseFrame(robot_state, this->m_cfg.tcp_offset); Eigen::Affine3d T_EE_in_base_frame = T_EE_in_base_frame_pose.affine_matrix(); @@ -495,8 +516,6 @@ void Franka::osc() { std::array tau_d_rate_limited = franka::limitRate( franka::kMaxTorqueRate, tau_d_array, robot_state.tau_J_d); - // deoxys/config/control_config.yml - std::array torque_limit = {5, 5, 5, 5, 5, 5, 5}; TorqueSafetyGuardFn(tau_d_rate_limited, torque_limit); return tau_d_rate_limited; @@ -514,17 +533,21 @@ void Franka::joint_controller() { franka::Model model = this->robot.loadModel(); const common::Vector7d Kp = this->m_cfg.kp; const common::Vector7d Kd = this->m_cfg.kd; + const common::Vector7d torque_limit = this->m_cfg.torque_limit; + const bool allow_high_collision = this->m_cfg.allow_high_collision; this->controller_time = 0.0; // conservative collision and impedance behavior this->set_default_robot_behavior(); - // high collision threshold values for high impedance - this->robot.setCollisionBehavior( - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + if (allow_high_collision) { + // High collision threshold values for high impedance. + this->robot.setCollisionBehavior( + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + } Eigen::Array joint_max_; Eigen::Array joint_min_; @@ -589,8 +612,6 @@ void Franka::joint_controller() { std::array tau_d_rate_limited = franka::limitRate( franka::kMaxTorqueRate, tau_d_array, robot_state.tau_J_d); - // deoxys/config/control_config.yml - std::array torque_limit = {5, 5, 5, 5, 5, 5, 5}; TorqueSafetyGuardFn(tau_d_rate_limited, torque_limit); return tau_d_rate_limited; @@ -615,12 +636,15 @@ void Franka::zero_torque_guiding() { } void Franka::zero_torque_controller() { - // high collision threshold values for high impedance - robot.setCollisionBehavior( - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, - {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + this->set_default_robot_behavior(); + if (this->m_cfg.allow_high_collision) { + // High collision threshold values for high impedance. + robot.setCollisionBehavior( + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}, + {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}}); + } this->controller_time = 0.0; try { @@ -730,34 +754,19 @@ std::optional> Franka::get_ik() { void Franka::set_cartesian_position(const common::Pose& x) { // pose is assumed to be in the robots coordinate frame - common::Pose target_pose = x; - if (!this->m_cfg.tcp_offset_configured_in_desk) { - target_pose = target_pose * this->m_cfg.tcp_offset.inverse(); - } if (this->m_cfg.async_control) { - this->osc_set_cartesian_position(target_pose); + this->osc_set_cartesian_position(x); return; } - // TODO: this should handled with tcp offset config - common::Pose nominal_end_effector_frame_value; - if (this->m_cfg.nominal_end_effector_frame.has_value()) { - nominal_end_effector_frame_value = - this->m_cfg.nominal_end_effector_frame.value(); - } else { - nominal_end_effector_frame_value = common::Pose::Identity(); - } - // nominal end effector frame should be on top of tcp offset as franka already - // takes care of the default franka hand offset lets add a franka hand offset - if (this->m_cfg.ik_solver == IKSolver::franka_ik) { - // if gripper is attached the tcp offset will automatically be applied - // by libfranka - this->robot.setEE(nominal_end_effector_frame_value.affine_array()); + const franka::RobotState robot_state = this->robot.readOnce(); + const common::Pose target_pose = + x * this->m_cfg.tcp_offset.inverse() * common::Pose(robot_state.F_T_EE); this->set_cartesian_position_internal(target_pose, 1.0, std::nullopt, std::nullopt); } else if (this->m_cfg.ik_solver == IKSolver::rcs_ik) { - this->set_cartesian_position_ik(target_pose); + this->set_cartesian_position_ik(x); } } @@ -814,10 +823,7 @@ void Franka::set_cartesian_position_internal(const common::Pose& pose, if (time == 0) { initial_elbow = state.elbow_c; - initial_pose = - this->m_cfg.tcp_offset_configured_in_desk - ? common::Pose(state.O_T_EE) - : common::Pose(state.O_T_EE) * this->m_cfg.tcp_offset; + initial_pose = common::Pose(state.O_T_EE); } auto new_elbow = initial_elbow; const double progress = time / max_time; diff --git a/extensions/rcs_fr3/src/hw/Franka.h b/extensions/rcs_fr3/src/hw/Franka.h index 6c66ba38..3906f49f 100644 --- a/extensions/rcs_fr3/src/hw/Franka.h +++ b/extensions/rcs_fr3/src/hw/Franka.h @@ -38,19 +38,22 @@ struct FrankaConfig : common::RobotConfig { common::RobotPlatform robot_platform = common::RobotPlatform::HARDWARE; IKSolver ik_solver = IKSolver::rcs_ik; double speed_factor = DEFAULT_SPEED_FACTOR; - // deoxys/config/joint-impedance-controller.yml + // values from deoxys/config/joint-impedance-controller.yml common::Vector7d kp = (common::Vector7d() << 100., 100., 100., 100., 75., 150., 50.).finished(); common::Vector7d kd = (common::Vector7d() << 20., 20., 20., 20., 7.5, 15.0, 5.0).finished(); + common::Vector7d torque_limit = common::Vector7d::Constant(5.0); + bool allow_high_collision = false; // values from deoxys/config/osc-position-controller.yml Eigen::Vector3d kp_p = (Eigen::Vector3d() << 150., 150., 150.).finished(); double kp_r = 250.0; std::optional load_parameters = std::nullopt; - std::optional nominal_end_effector_frame = std::nullopt; std::optional world_to_robot = std::nullopt; + common::Pose tcp_offset = common::Pose::Identity(); + // Indicates that Cartesian control uses tcp_offset. + bool tcp_offset_explicit = false; bool async_control = false; - bool tcp_offset_configured_in_desk = true; bool ignore_realtime = false; size_t dof = 7; Eigen::Matrix joint_limits = @@ -118,6 +121,8 @@ class Franka : public common::Robot { common::Pose get_cartesian_position() override; + common::Pose get_cartesian_flange_position() override; + void set_joint_position(const common::VectorXd& q) override; common::VectorXd get_joint_position() override; diff --git a/extensions/rcs_fr3/src/pybind/rcs.cpp b/extensions/rcs_fr3/src/pybind/rcs.cpp index 51c8a4a0..9b2f4ece 100644 --- a/extensions/rcs_fr3/src/pybind/rcs.cpp +++ b/extensions/rcs_fr3/src/pybind/rcs.cpp @@ -130,122 +130,129 @@ PYBIND11_MODULE(_core, m) { .def_readwrite("speed_factor", &rcs::hw::FrankaConfig::speed_factor) .def_readwrite("kp", &rcs::hw::FrankaConfig::kp) .def_readwrite("kd", &rcs::hw::FrankaConfig::kd) + .def_readwrite("torque_limit", &rcs::hw::FrankaConfig::torque_limit) + .def_readwrite("allow_high_collision", + &rcs::hw::FrankaConfig::allow_high_collision) .def_readwrite("kp_p", &rcs::hw::FrankaConfig::kp_p) .def_readwrite("kp_r", &rcs::hw::FrankaConfig::kp_r) .def_readwrite("load_parameters", &rcs::hw::FrankaConfig::load_parameters) - .def_readwrite("nominal_end_effector_frame", - &rcs::hw::FrankaConfig::nominal_end_effector_frame) .def_readwrite("world_to_robot", &rcs::hw::FrankaConfig::world_to_robot) - .def_readwrite("tcp_offset_configured_in_desk", - &rcs::hw::FrankaConfig::tcp_offset_configured_in_desk) + .def_property( + "tcp_offset", + [](const rcs::hw::FrankaConfig& config) { return config.tcp_offset; }, + [](rcs::hw::FrankaConfig& config, + const rcs::common::Pose& tcp_offset) { + config.tcp_offset = tcp_offset; + config.tcp_offset_explicit = true; + }) + .def_readwrite("tcp_offset_explicit", + &rcs::hw::FrankaConfig::tcp_offset_explicit) .def_readwrite("async_control", &rcs::hw::FrankaConfig::async_control) .def_readwrite("ignore_realtime", &rcs::hw::FrankaConfig::ignore_realtime) .def_readwrite("ip", &rcs::hw::FrankaConfig::ip); rcs::hw::FR3Config default_fr3_config; py::class_(hw, "FR3Config") - .def(py::init( - [](const std::string& ip, rcs::hw::IKSolver ik_solver, - double speed_factor, const rcs::common::Vector7d& kp, - const rcs::common::Vector7d& kd, const Eigen::Vector3d& kp_p, - double kp_r, - std::optional load_parameters, - std::optional nominal_end_effector_frame, - std::optional world_to_robot, - bool async_control, bool tcp_offset_configured_in_desk, - bool ignore_realtime, rcs::common::Pose tcp_offset, - std::string attachment_site, - std::string kinematic_model_path) { - rcs::hw::FR3Config cfg; - cfg.ik_solver = ik_solver; - cfg.speed_factor = speed_factor; - cfg.kp = kp; - cfg.kd = kd; - cfg.kp_p = kp_p; - cfg.kp_r = kp_r; - cfg.load_parameters = load_parameters; - cfg.nominal_end_effector_frame = nominal_end_effector_frame; - cfg.world_to_robot = world_to_robot; - cfg.async_control = async_control; - cfg.tcp_offset_configured_in_desk = - tcp_offset_configured_in_desk; - cfg.ignore_realtime = ignore_realtime; - cfg.ip = ip; - cfg.tcp_offset = tcp_offset; - cfg.attachment_site = attachment_site; - cfg.kinematic_model_path = kinematic_model_path; - return cfg; - }), - py::arg("ip"), py::arg("ik_solver") = default_fr3_config.ik_solver, - py::arg("speed_factor") = default_fr3_config.speed_factor, - py::arg("kp") = default_fr3_config.kp, - py::arg("kd") = default_fr3_config.kd, - py::arg("kp_p") = default_fr3_config.kp_p, - py::arg("kp_r") = default_fr3_config.kp_r, - py::arg("load_parameters") = default_fr3_config.load_parameters, - py::arg("nominal_end_effector_frame") = - default_fr3_config.nominal_end_effector_frame, - py::arg("world_to_robot") = default_fr3_config.world_to_robot, - py::arg("async_control") = default_fr3_config.async_control, - py::arg("tcp_offset_configured_in_desk") = - default_fr3_config.tcp_offset_configured_in_desk, - py::arg("ignore_realtime") = default_fr3_config.ignore_realtime, - py::arg("tcp_offset") = default_fr3_config.tcp_offset, - py::arg("attachment_site") = default_fr3_config.attachment_site, - py::arg("kinematic_model_path") = - default_fr3_config.kinematic_model_path); + .def( + py::init([](const std::string& ip, rcs::hw::IKSolver ik_solver, + double speed_factor, const rcs::common::Vector7d& kp, + const rcs::common::Vector7d& kd, + const Eigen::Vector3d& kp_p, double kp_r, + std::optional load_parameters, + std::optional world_to_robot, + bool async_control, bool ignore_realtime, + rcs::common::Pose tcp_offset, std::string attachment_site, + std::string kinematic_model_path, + const rcs::common::Vector7d& torque_limit, + bool allow_high_collision) { + rcs::hw::FR3Config cfg; + cfg.ik_solver = ik_solver; + cfg.speed_factor = speed_factor; + cfg.kp = kp; + cfg.kd = kd; + cfg.torque_limit = torque_limit; + cfg.allow_high_collision = allow_high_collision; + cfg.kp_p = kp_p; + cfg.kp_r = kp_r; + cfg.load_parameters = load_parameters; + cfg.world_to_robot = world_to_robot; + cfg.async_control = async_control; + cfg.ignore_realtime = ignore_realtime; + cfg.ip = ip; + cfg.tcp_offset = tcp_offset; + cfg.tcp_offset_explicit = true; + cfg.attachment_site = attachment_site; + cfg.kinematic_model_path = kinematic_model_path; + return cfg; + }), + py::arg("ip"), py::arg("ik_solver") = default_fr3_config.ik_solver, + py::arg("speed_factor") = default_fr3_config.speed_factor, + py::arg("kp") = default_fr3_config.kp, + py::arg("kd") = default_fr3_config.kd, + py::arg("kp_p") = default_fr3_config.kp_p, + py::arg("kp_r") = default_fr3_config.kp_r, + py::arg("load_parameters") = default_fr3_config.load_parameters, + py::arg("world_to_robot") = default_fr3_config.world_to_robot, + py::arg("async_control") = default_fr3_config.async_control, + py::arg("ignore_realtime") = default_fr3_config.ignore_realtime, + py::arg("tcp_offset") = default_fr3_config.tcp_offset, + py::arg("attachment_site") = default_fr3_config.attachment_site, + py::arg("kinematic_model_path") = + default_fr3_config.kinematic_model_path, + py::arg("torque_limit") = default_fr3_config.torque_limit, + py::arg("allow_high_collision") = + default_fr3_config.allow_high_collision); rcs::hw::PandaConfig default_panda_config; py::class_(hw, "PandaConfig") - .def(py::init( - [](const std::string& ip, rcs::hw::IKSolver ik_solver, - double speed_factor, const rcs::common::Vector7d& kp, - const rcs::common::Vector7d& kd, const Eigen::Vector3d& kp_p, - double kp_r, - std::optional load_parameters, - std::optional nominal_end_effector_frame, - std::optional world_to_robot, - bool async_control, bool tcp_offset_configured_in_desk, - bool ignore_realtime, rcs::common::Pose tcp_offset, - std::string attachment_site, - std::string kinematic_model_path) { - rcs::hw::PandaConfig cfg; - cfg.ik_solver = ik_solver; - cfg.speed_factor = speed_factor; - cfg.kp = kp; - cfg.kd = kd; - cfg.kp_p = kp_p; - cfg.kp_r = kp_r; - cfg.load_parameters = load_parameters; - cfg.nominal_end_effector_frame = nominal_end_effector_frame; - cfg.world_to_robot = world_to_robot; - cfg.async_control = async_control; - cfg.tcp_offset_configured_in_desk = - tcp_offset_configured_in_desk; - cfg.ignore_realtime = ignore_realtime; - cfg.ip = ip; - cfg.tcp_offset = tcp_offset; - cfg.attachment_site = attachment_site; - cfg.kinematic_model_path = kinematic_model_path; - return cfg; - }), - py::arg("ip"), py::arg("ik_solver") = default_panda_config.ik_solver, - py::arg("speed_factor") = default_panda_config.speed_factor, - py::arg("kp") = default_panda_config.kp, - py::arg("kd") = default_panda_config.kd, - py::arg("kp_p") = default_panda_config.kp_p, - py::arg("kp_r") = default_panda_config.kp_r, - py::arg("load_parameters") = default_panda_config.load_parameters, - py::arg("nominal_end_effector_frame") = - default_panda_config.nominal_end_effector_frame, - py::arg("world_to_robot") = default_panda_config.world_to_robot, - py::arg("async_control") = default_panda_config.async_control, - py::arg("tcp_offset_configured_in_desk") = - default_panda_config.tcp_offset_configured_in_desk, - py::arg("ignore_realtime") = default_panda_config.ignore_realtime, - py::arg("tcp_offset") = default_panda_config.tcp_offset, - py::arg("attachment_site") = default_panda_config.attachment_site, - py::arg("kinematic_model_path") = - default_panda_config.kinematic_model_path); + .def( + py::init([](const std::string& ip, rcs::hw::IKSolver ik_solver, + double speed_factor, const rcs::common::Vector7d& kp, + const rcs::common::Vector7d& kd, + const Eigen::Vector3d& kp_p, double kp_r, + std::optional load_parameters, + std::optional world_to_robot, + bool async_control, bool ignore_realtime, + rcs::common::Pose tcp_offset, std::string attachment_site, + std::string kinematic_model_path, + const rcs::common::Vector7d& torque_limit, + bool allow_high_collision) { + rcs::hw::PandaConfig cfg; + cfg.ik_solver = ik_solver; + cfg.speed_factor = speed_factor; + cfg.kp = kp; + cfg.kd = kd; + cfg.torque_limit = torque_limit; + cfg.allow_high_collision = allow_high_collision; + cfg.kp_p = kp_p; + cfg.kp_r = kp_r; + cfg.load_parameters = load_parameters; + cfg.world_to_robot = world_to_robot; + cfg.async_control = async_control; + cfg.ignore_realtime = ignore_realtime; + cfg.ip = ip; + cfg.tcp_offset = tcp_offset; + cfg.tcp_offset_explicit = true; + cfg.attachment_site = attachment_site; + cfg.kinematic_model_path = kinematic_model_path; + return cfg; + }), + py::arg("ip"), py::arg("ik_solver") = default_panda_config.ik_solver, + py::arg("speed_factor") = default_panda_config.speed_factor, + py::arg("kp") = default_panda_config.kp, + py::arg("kd") = default_panda_config.kd, + py::arg("kp_p") = default_panda_config.kp_p, + py::arg("kp_r") = default_panda_config.kp_r, + py::arg("load_parameters") = default_panda_config.load_parameters, + py::arg("world_to_robot") = default_panda_config.world_to_robot, + py::arg("async_control") = default_panda_config.async_control, + py::arg("ignore_realtime") = default_panda_config.ignore_realtime, + py::arg("tcp_offset") = default_panda_config.tcp_offset, + py::arg("attachment_site") = default_panda_config.attachment_site, + py::arg("kinematic_model_path") = + default_panda_config.kinematic_model_path, + py::arg("torque_limit") = default_panda_config.torque_limit, + py::arg("allow_high_collision") = + default_panda_config.allow_high_collision); py::object gripper_config = (py::object)py::module_::import("rcs").attr("common").attr( @@ -305,6 +312,8 @@ PYBIND11_MODULE(_core, m) { .def("set_config", &rcs::hw::Franka::set_config, py::arg("cfg")) .def("get_config", &rcs::hw::Franka::get_config) .def("get_state", &rcs::hw::Franka::get_state) + .def("get_cartesian_flange_position", + &rcs::hw::Franka::get_cartesian_flange_position) .def("set_default_robot_behavior", &rcs::hw::Franka::set_default_robot_behavior) .def("set_guiding_mode", &rcs::hw::Franka::set_guiding_mode, diff --git a/extensions/rcs_fr3/src/rcs_fr3/_core/hw/__init__.pyi b/extensions/rcs_fr3/src/rcs_fr3/_core/hw/__init__.pyi index 84947baa..af761256 100644 --- a/extensions/rcs_fr3/src/rcs_fr3/_core/hw/__init__.pyi +++ b/extensions/rcs_fr3/src/rcs_fr3/_core/hw/__init__.pyi @@ -80,6 +80,7 @@ class Franka(rcs._core.common.Robot): self, desired_q: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] ) -> None: ... def double_tap_robot_to_continue(self) -> None: ... + def get_cartesian_flange_position(self) -> rcs._core.common.Pose: ... def get_config(self) -> FrankaConfig: ... def get_state(self) -> FrankaState: ... def osc_set_cartesian_position(self, desired_pos_EE_in_base_frame: rcs._core.common.Pose) -> None: ... @@ -103,6 +104,7 @@ class Franka(rcs._core.common.Robot): def zero_torque_guiding(self) -> None: ... class FrankaConfig(rcs._core.common.RobotConfig): + allow_high_collision: bool async_control: bool ignore_realtime: bool ik_solver: IKSolver @@ -112,9 +114,10 @@ class FrankaConfig(rcs._core.common.RobotConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] kp_r: float load_parameters: FrankaLoad | None - nominal_end_effector_frame: rcs._core.common.Pose | None speed_factor: float - tcp_offset_configured_in_desk: bool + tcp_offset: rcs._core.common.Pose + tcp_offset_explicit: bool + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] world_to_robot: rcs._core.common.Pose | None class FrankaHand(rcs._core.common.Gripper): @@ -307,14 +310,14 @@ class FR3Config(FrankaConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] = ..., kp_r: float = 250.0, load_parameters: FrankaLoad | None = None, - nominal_end_effector_frame: rcs._core.common.Pose | None = None, world_to_robot: rcs._core.common.Pose | None = None, async_control: bool = False, - tcp_offset_configured_in_desk: bool = True, ignore_realtime: bool = False, tcp_offset: rcs._core.common.Pose = ..., attachment_site: str = "attachment_site", kinematic_model_path: str = "assets/scenes/fr3_empty_world/robot.xml", + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ..., + allow_high_collision: bool = False, ) -> None: ... class PandaConfig(FrankaConfig): @@ -328,14 +331,14 @@ class PandaConfig(FrankaConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] = ..., kp_r: float = 250.0, load_parameters: FrankaLoad | None = None, - nominal_end_effector_frame: rcs._core.common.Pose | None = None, world_to_robot: rcs._core.common.Pose | None = None, async_control: bool = False, - tcp_offset_configured_in_desk: bool = True, ignore_realtime: bool = False, tcp_offset: rcs._core.common.Pose = ..., attachment_site: str = "attachment_site", kinematic_model_path: str = "assets/scenes/fr3_empty_world/robot.xml", + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ..., + allow_high_collision: bool = False, ) -> None: ... class FrankaState(rcs._core.common.RobotState): diff --git a/extensions/rcs_fr3/src/rcs_fr3/configs.py b/extensions/rcs_fr3/src/rcs_fr3/configs.py index 2a25d056..c989bc2c 100644 --- a/extensions/rcs_fr3/src/rcs_fr3/configs.py +++ b/extensions/rcs_fr3/src/rcs_fr3/configs.py @@ -25,7 +25,6 @@ def config(self) -> FR3HardwareEnvCreatorConfig: ik_solver=hw.IKSolver.rcs_ik, speed_factor=0.1, async_control=False, - tcp_offset_configured_in_desk=True, ignore_realtime=False, tcp_offset=common.Pose(common.FrankaHandTCPOffset()), attachment_site="attachment_site", diff --git a/extensions/rcs_fr3/src/rcs_fr3/desk.py b/extensions/rcs_fr3/src/rcs_fr3/desk.py index 32d6c0ed..10006252 100644 --- a/extensions/rcs_fr3/src/rcs_fr3/desk.py +++ b/extensions/rcs_fr3/src/rcs_fr3/desk.py @@ -51,6 +51,7 @@ def home(ip: str): env_cfg = default_env.config() robot_cfg = env_cfg.robot_cfg robot_cfg.speed_factor = 0.2 + robot_cfg.ignore_realtime = True f = rcs_fr3.hw.Franka(robot_cfg) f.move_home() @@ -72,6 +73,7 @@ def gripper(ip: str, close_gripper: bool): def info(ip: str, include_hand: bool = False): robot_cfg = rcs_fr3.hw.FR3Config(ip=ip) robot_cfg.speed_factor = 0.2 + robot_cfg.ignore_realtime = True f = rcs_fr3.hw.Franka(robot_cfg) print("Robot info:") print("Current cartesian position:") diff --git a/extensions/rcs_fr3/src/rcs_fr3/envs.py b/extensions/rcs_fr3/src/rcs_fr3/envs.py index 8b780691..3bd10c58 100644 --- a/extensions/rcs_fr3/src/rcs_fr3/envs.py +++ b/extensions/rcs_fr3/src/rcs_fr3/envs.py @@ -44,7 +44,8 @@ def _rs2dict(self, state: hw.RobotState): def reset( self, *, seed: int | None = None, options: dict[str, Any] | None = None ) -> tuple[dict[str, Any], dict[str, Any]]: - return super().reset(seed=seed, options=options) + obs, info = super().reset(seed=seed, options=options) + return self.get_obs(obs), info def close(self): self.hw_robot.stop_control_thread() diff --git a/extensions/rcs_panda/src/rcs_panda/_core/hw/__init__.pyi b/extensions/rcs_panda/src/rcs_panda/_core/hw/__init__.pyi index 84947baa..af761256 100644 --- a/extensions/rcs_panda/src/rcs_panda/_core/hw/__init__.pyi +++ b/extensions/rcs_panda/src/rcs_panda/_core/hw/__init__.pyi @@ -80,6 +80,7 @@ class Franka(rcs._core.common.Robot): self, desired_q: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] ) -> None: ... def double_tap_robot_to_continue(self) -> None: ... + def get_cartesian_flange_position(self) -> rcs._core.common.Pose: ... def get_config(self) -> FrankaConfig: ... def get_state(self) -> FrankaState: ... def osc_set_cartesian_position(self, desired_pos_EE_in_base_frame: rcs._core.common.Pose) -> None: ... @@ -103,6 +104,7 @@ class Franka(rcs._core.common.Robot): def zero_torque_guiding(self) -> None: ... class FrankaConfig(rcs._core.common.RobotConfig): + allow_high_collision: bool async_control: bool ignore_realtime: bool ik_solver: IKSolver @@ -112,9 +114,10 @@ class FrankaConfig(rcs._core.common.RobotConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] kp_r: float load_parameters: FrankaLoad | None - nominal_end_effector_frame: rcs._core.common.Pose | None speed_factor: float - tcp_offset_configured_in_desk: bool + tcp_offset: rcs._core.common.Pose + tcp_offset_explicit: bool + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] world_to_robot: rcs._core.common.Pose | None class FrankaHand(rcs._core.common.Gripper): @@ -307,14 +310,14 @@ class FR3Config(FrankaConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] = ..., kp_r: float = 250.0, load_parameters: FrankaLoad | None = None, - nominal_end_effector_frame: rcs._core.common.Pose | None = None, world_to_robot: rcs._core.common.Pose | None = None, async_control: bool = False, - tcp_offset_configured_in_desk: bool = True, ignore_realtime: bool = False, tcp_offset: rcs._core.common.Pose = ..., attachment_site: str = "attachment_site", kinematic_model_path: str = "assets/scenes/fr3_empty_world/robot.xml", + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ..., + allow_high_collision: bool = False, ) -> None: ... class PandaConfig(FrankaConfig): @@ -328,14 +331,14 @@ class PandaConfig(FrankaConfig): kp_p: numpy.ndarray[tuple[typing.Literal[3]], numpy.dtype[numpy.float64]] = ..., kp_r: float = 250.0, load_parameters: FrankaLoad | None = None, - nominal_end_effector_frame: rcs._core.common.Pose | None = None, world_to_robot: rcs._core.common.Pose | None = None, async_control: bool = False, - tcp_offset_configured_in_desk: bool = True, ignore_realtime: bool = False, tcp_offset: rcs._core.common.Pose = ..., attachment_site: str = "attachment_site", kinematic_model_path: str = "assets/scenes/fr3_empty_world/robot.xml", + torque_limit: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ..., + allow_high_collision: bool = False, ) -> None: ... class FrankaState(rcs._core.common.RobotState): diff --git a/extensions/rcs_panda/src/rcs_panda/desk.py b/extensions/rcs_panda/src/rcs_panda/desk.py index f64f8d96..5ca6bd4d 100644 --- a/extensions/rcs_panda/src/rcs_panda/desk.py +++ b/extensions/rcs_panda/src/rcs_panda/desk.py @@ -51,6 +51,7 @@ def home(ip: str): env_cfg = default_env.config() robot_cfg = env_cfg.robot_cfg robot_cfg.speed_factor = 0.2 + robot_cfg.ignore_realtime = True f = rcs_panda.hw.Franka(robot_cfg) f.move_home() @@ -71,6 +72,7 @@ def gripper(ip: str, close_gripper: bool): def info(ip: str, include_hand: bool = False): robot_cfg = rcs_panda.hw.PandaConfig(ip=ip) robot_cfg.speed_factor = 0.2 + robot_cfg.ignore_realtime = True f = rcs_panda.hw.Franka(robot_cfg) print("Robot info:") print("Current cartesian position:") diff --git a/extensions/rcs_panda/src/rcs_panda/envs.py b/extensions/rcs_panda/src/rcs_panda/envs.py index 34e8d1b1..2e83794f 100644 --- a/extensions/rcs_panda/src/rcs_panda/envs.py +++ b/extensions/rcs_panda/src/rcs_panda/envs.py @@ -44,7 +44,8 @@ def _rs2dict(self, state: hw.RobotState): def reset( self, *, seed: int | None = None, options: dict[str, Any] | None = None ) -> tuple[dict[str, Any], dict[str, Any]]: - return super().reset(seed=seed, options=options) + obs, info = super().reset(seed=seed, options=options) + return self.get_obs(obs), info def close(self): self.hw_robot.stop_control_thread() diff --git a/extensions/rcs_so101/src/rcs_so101/hw.py b/extensions/rcs_so101/src/rcs_so101/hw.py index d3e50718..3c44dee6 100644 --- a/extensions/rcs_so101/src/rcs_so101/hw.py +++ b/extensions/rcs_so101/src/rcs_so101/hw.py @@ -50,6 +50,9 @@ def __init__(self, cfg: SO101Config, ik: common.Kinematics): def get_cartesian_position(self) -> common.Pose: return self.ik.forward(self.get_joint_position()) + def get_cartesian_flange_position(self) -> common.Pose: + return self.get_cartesian_position() + def get_ik(self) -> common.Kinematics | None: return self.ik diff --git a/extensions/rcs_ur5e/src/rcs_ur5e/hw.py b/extensions/rcs_ur5e/src/rcs_ur5e/hw.py index 56b371e2..9fa9530a 100644 --- a/extensions/rcs_ur5e/src/rcs_ur5e/hw.py +++ b/extensions/rcs_ur5e/src/rcs_ur5e/hw.py @@ -269,6 +269,9 @@ def get_cartesian_position(self) -> common.Pose: common.Pose(rpy_vector=np.array([0, 0, np.deg2rad(180)]), translation=np.array([0, 0, 0])).inverse() * pose # type: ignore ) + def get_cartesian_flange_position(self) -> common.Pose: + return self.get_cartesian_position() + def get_ik(self) -> common.Kinematics | None: return self.ik diff --git a/extensions/rcs_xarm7/src/rcs_xarm7/hw.py b/extensions/rcs_xarm7/src/rcs_xarm7/hw.py index d5ece995..dd48f56d 100644 --- a/extensions/rcs_xarm7/src/rcs_xarm7/hw.py +++ b/extensions/rcs_xarm7/src/rcs_xarm7/hw.py @@ -60,6 +60,9 @@ def get_cartesian_position(self) -> common.Pose: return common.Pose(rpy_vector=rpy, translation=translation_meter) # type: ignore + def get_cartesian_flange_position(self) -> common.Pose: + return self.get_cartesian_position() + def get_ik(self) -> common.Kinematics | None: return self.ik diff --git a/extensions/rcs_yam/src/rcs_yam/hw.py b/extensions/rcs_yam/src/rcs_yam/hw.py index 76cd0e5d..24063733 100644 --- a/extensions/rcs_yam/src/rcs_yam/hw.py +++ b/extensions/rcs_yam/src/rcs_yam/hw.py @@ -123,8 +123,10 @@ def get_gripper_width(self) -> float: def get_cartesian_position(self) -> common.Pose: # `Kinematics.forward` applies the inverse of the offset it is handed, so the TCP is composed # here instead, to match the pose `SimRobot::get_cartesian_position` reports in simulation. - flange = self.ik.forward(self.get_joint_position(), common.Pose()) - return flange * self._config.tcp_offset + return self.get_cartesian_flange_position() * self._config.tcp_offset + + def get_cartesian_flange_position(self) -> common.Pose: + return self.ik.forward(self.get_joint_position(), common.Pose()) def set_cartesian_position(self, pose: common.Pose) -> None: q = self.ik.inverse(pose, self.get_joint_position(), self._config.tcp_offset) diff --git a/extensions/rcs_zed/src/rcs_zed/__main__.py b/extensions/rcs_zed/src/rcs_zed/__main__.py index ba6c3771..7e184b82 100644 --- a/extensions/rcs_zed/src/rcs_zed/__main__.py +++ b/extensions/rcs_zed/src/rcs_zed/__main__.py @@ -62,6 +62,52 @@ def serials(): typer.echo(f" {device.model}: {device.serial} (imu={device.has_imu})") +def _format_intrinsics(name: str, matrix) -> str: + fx, fy, cx, cy = float(matrix[0, 0]), float(matrix[1, 1]), float(matrix[0, 2]), float(matrix[1, 2]) + lines = [ + f"{name}:", + f" fx={fx:.6f} fy={fy:.6f} cx={cx:.6f} cy={cy:.6f}", + f" K=\n{matrix}", + ] + return "\n".join(lines) + + +@zed_app.command() +def intrinsics( + serial: str | None = typer.Argument(None, help="Optional ZED serial number. Uses the first device if omitted."), + width: int = typer.Option(1280, help="Requested capture width."), + height: int = typer.Option(720, help="Requested capture height."), + fps: int = typer.Option(30, help="Requested capture frame rate."), + right: bool = typer.Option(False, "--right", help="Also print right-camera intrinsics."), +): + """Open a ZED camera and print its RGB intrinsics for the requested resolution.""" + serial = _resolve_serial(serial) + config = common.BaseCameraConfig( + identifier=serial, + resolution_width=width, + resolution_height=height, + frame_rate=fps, + ) + try: + handle = ZEDCameraSet.open_camera( + config, + enable_depth=False, + enable_imu=False, + include_right=right, + ) + except Exception as exc: + msg = f"Could not start ZED camera {serial}: {exc}" + raise typer.BadParameter(msg) from exc + + try: + typer.echo(f"serial={handle.device_info.serial} model={handle.device_info.model} resolution={width}x{height}") + typer.echo(_format_intrinsics("left", handle.color_intrinsics)) + if right and handle.right_color_intrinsics is not None: + typer.echo(_format_intrinsics("right", handle.right_color_intrinsics)) + finally: + handle.close() + + @zed_app.command("rgb-view") def rgb_view( serial: str | None = typer.Argument(None, help="Optional ZED serial number. Uses the first device if omitted."), diff --git a/extensions/rcs_zed/src/rcs_zed/camera.py b/extensions/rcs_zed/src/rcs_zed/camera.py index 881a1368..8e86481d 100644 --- a/extensions/rcs_zed/src/rcs_zed/camera.py +++ b/extensions/rcs_zed/src/rcs_zed/camera.py @@ -216,7 +216,6 @@ def open_camera( init.camera_fps = config.frame_rate init.coordinate_units = sl.UNIT.METER init.depth_mode = sl.DEPTH_MODE.NONE if not enable_depth else sl.DEPTH_MODE.QUALITY - init.sdk_verbose = False init.set_from_serial_number(int(config.identifier)) camera = sl.Camera() diff --git a/include/rcs/Robot.h b/include/rcs/Robot.h index 9a94cf4a..e7b680c7 100644 --- a/include/rcs/Robot.h +++ b/include/rcs/Robot.h @@ -116,6 +116,10 @@ class Robot { virtual Pose get_cartesian_position() = 0; + virtual Pose get_cartesian_flange_position() { + return get_cartesian_position(); + } + virtual void set_joint_position(const VectorXd& q) = 0; virtual VectorXd get_joint_position() = 0; diff --git a/python/rcs/__main__.py b/python/rcs/__main__.py index 4f97eae0..0df6232f 100644 --- a/python/rcs/__main__.py +++ b/python/rcs/__main__.py @@ -21,7 +21,7 @@ run_conversion, ) from rcs.sim.replayer import replay as replay_dataset -from rcs.utils import export_episode_videos +from rcs.utils import export_camera_episode_videos, export_episode_videos app = typer.Typer() @@ -218,5 +218,44 @@ def episode_videos( export_episode_videos(dataset=dataset, output=output, fps=fps, n=n) +@app.command("camera-episode-videos") +def camera_episode_videos( + dataset: Annotated[ + Path, + typer.Argument( + exists=True, + help="Parquet dataset file or directory with parquet parts.", + ), + ], + output: Annotated[ + Path, + typer.Argument( + exists=False, + help="Output directory for camera episode mp4 files.", + ), + ], + fps: Annotated[int, typer.Option(help="Video frames per second.")] = DEFAULT_FPS, + camera: Annotated[ + str | None, + typer.Option(help="Only export this camera. By default, exports every camera."), + ] = None, + episode: Annotated[ + int | None, + typer.Option(help="Only export this zero-based, recording-order episode."), + ] = None, +): + """Export a simple raw-frame MP4 for every camera in every episode.""" + try: + export_camera_episode_videos( + dataset=dataset, + output=output, + fps=fps, + camera=camera, + episode=episode, + ) + except ValueError as error: + raise typer.BadParameter(str(error)) from error + + if __name__ == "__main__": app() diff --git a/python/rcs/_core/common.pyi b/python/rcs/_core/common.pyi index b2749517..95ec05fd 100644 --- a/python/rcs/_core/common.pyi +++ b/python/rcs/_core/common.pyi @@ -226,6 +226,7 @@ class Robot: def __init__(self) -> None: ... def close(self) -> None: ... def get_base_pose_in_world_coordinates(self) -> Pose: ... + def get_cartesian_flange_position(self) -> Pose: ... def get_cartesian_position(self) -> Pose: ... def get_config(self) -> RobotConfig: ... def get_ik(self) -> Kinematics | None: ... diff --git a/python/rcs/envs/base.py b/python/rcs/envs/base.py index c5d8da0e..62bc897f 100644 --- a/python/rcs/envs/base.py +++ b/python/rcs/envs/base.py @@ -167,7 +167,15 @@ class CameraDictType(RCSpaceType): # joining works with inheritance but need to inherit from protocol again -class ArmObsType(TQuatDictType, JointsDictType, TRPYDictType): ... +class ArmObsType(TQuatDictType, JointsDictType, TRPYDictType): + tquat_flange: Annotated[ + Vec7Type, + gym.spaces.Box( + low=np.array([-0.855, -0.855, -1] + [-1] + [-np.inf] * 3), + high=np.array([0.855, 0.855, 1.188] + [1] + [np.inf] * 3), + dtype=np.float64, + ), + ] CartOrJointContType: TypeAlias = TQuatDictType | JointsDictType | TRPYDictType @@ -346,6 +354,12 @@ def get_robot_obs(self) -> ArmObsType: ), joints=self.robot.get_joint_position(), xyzrpy=self.robot.get_cartesian_position().xyzrpy(), + tquat_flange=np.concatenate( + [ + self.robot.get_cartesian_flange_position().translation(), + self.robot.get_cartesian_flange_position().rotation_q(), + ] + ), ) def action(self, action: dict[str, Any]) -> dict[str, Any]: diff --git a/python/rcs/utils.py b/python/rcs/utils.py index 5f5aa372..57a80b6e 100644 --- a/python/rcs/utils.py +++ b/python/rcs/utils.py @@ -133,6 +133,24 @@ def _render_action_panel( return image +def _episode_starts(conn: duckdb.DuckDBPyConnection, source_escaped: str) -> list[tuple[str, float]]: + """Return episodes in recording order, with a stable tie-breaker.""" + return conn.execute( + f""" + SELECT uuid, MIN(timestamp) AS start_timestamp + FROM read_parquet('{source_escaped}') + GROUP BY uuid + ORDER BY start_timestamp, uuid + """ + ).fetchall() + + +def _episode_filename(timestamp: float, episode_number: int, camera_name: str | None = None) -> str: + timestamp_text = datetime.datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d-%H-%M-%S") + filename = f"{timestamp_text}_episode-{episode_number:06d}" + return f"{filename}_{camera_name}.mp4" if camera_name is not None else f"{filename}.mp4" + + def export_episode_videos( dataset: str | Path, output: str | Path, @@ -162,8 +180,8 @@ def export_episode_videos( for robot, robot_struct in action_struct.children } - uuids = conn.execute(f"SELECT DISTINCT uuid FROM read_parquet('{source_escaped}') ORDER BY uuid").fetchall() - for index, (episode_id,) in enumerate(uuids): + episodes = _episode_starts(conn, source_escaped) + for index, (episode_id, _) in enumerate(episodes): if n != -1 and index >= n: break @@ -204,7 +222,6 @@ def export_episode_videos( if not rows: continue - timestamp = datetime.datetime.fromtimestamp(float(rows[0][0])).strftime("%Y-%m-%d-%H-%M-%S") frames = [] joint_history = { robot: np.asarray([row[1 + len(camera_names) + robot_idx] for row in rows], dtype=np.float32) @@ -233,4 +250,70 @@ def export_episode_videos( tiled[top : top + height, left : left + width] = image frames.append(tiled) - _write_mp4(frames, output / f"{timestamp}.mp4", fps=fps) + _write_mp4(frames, output / _episode_filename(float(rows[0][0]), index), fps=fps) + + +def export_camera_episode_videos( + dataset: str | Path, + output: str | Path, + fps: int = 30, + camera: str | None = None, + episode: int | None = None, +) -> None: + """Export raw camera frames as one MP4 for every selected camera and episode. + + ``episode`` is the zero-based recording-order index used in the filenames. + """ + import torch + from torchvision.io import decode_jpeg + + dataset = Path(dataset) + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + + source = str(dataset / "*.parquet") if dataset.is_dir() else str(dataset) + source_escaped = source.replace("'", "''") + conn = duckdb.connect() + relation = conn.sql(f"SELECT * FROM read_parquet('{source_escaped}')") + frame_struct = relation.select("obs.frames").types[0] + camera_names = [name for name, _ in frame_struct.children] + if camera is not None: + if camera not in camera_names: + available = ", ".join(camera_names) + message = f"Unknown camera {camera!r}. Available cameras: {available}" + raise ValueError(message) + camera_names = [camera] + + episodes = _episode_starts(conn, source_escaped) + if episode is not None: + if episode < 0 or episode >= len(episodes): + message = f"Episode {episode} is out of range (dataset has {len(episodes)} episodes)." + raise ValueError(message) + selected_episodes = [(episode, episodes[episode])] + else: + selected_episodes = list(enumerate(episodes)) + + for episode_number, (episode_id, _) in selected_episodes: + for camera_name in camera_names: + rows = conn.execute( + f""" + SELECT timestamp, obs.frames.{camera_name}.rgb.data + FROM read_parquet('{source_escaped}') + WHERE uuid = ? + AND obs.frames.{camera_name}.rgb.data IS NOT NULL + ORDER BY step + """, + [episode_id], + ).fetchall() + if not rows: + continue + + frames = [ + decode_jpeg(torch.frombuffer(bytearray(image_bytes), dtype=torch.uint8)).permute(1, 2, 0).cpu().numpy() + for _, image_bytes in rows + ] + _write_mp4( + frames, + output / _episode_filename(float(rows[0][0]), episode_number, camera_name), + fps=fps, + ) diff --git a/src/pybind/rcs.cpp b/src/pybind/rcs.cpp index 9be15edf..1f172b81 100644 --- a/src/pybind/rcs.cpp +++ b/src/pybind/rcs.cpp @@ -73,6 +73,11 @@ class PyRobot : public rcs::common::Robot { get_cartesian_position, ); } + rcs::common::Pose get_cartesian_flange_position() override { + PYBIND11_OVERRIDE(rcs::common::Pose, rcs::common::Robot, + get_cartesian_flange_position, ); + } + void set_joint_position(const rcs::common::VectorXd& q) override { PYBIND11_OVERRIDE_PURE(void, rcs::common::Robot, set_joint_position, q); } @@ -461,6 +466,8 @@ PYBIND11_MODULE(_core, m) { .def("get_state", &rcs::common::Robot::get_state) .def("get_cartesian_position", &rcs::common::Robot::get_cartesian_position) + .def("get_cartesian_flange_position", + &rcs::common::Robot::get_cartesian_flange_position) .def("set_joint_position", &rcs::common::Robot::set_joint_position, py::arg("q"), py::call_guard()) .def("get_joint_position", &rcs::common::Robot::get_joint_position) diff --git a/src/sim/SimRobot.cpp b/src/sim/SimRobot.cpp index c1263580..646a2387 100644 --- a/src/sim/SimRobot.cpp +++ b/src/sim/SimRobot.cpp @@ -117,12 +117,16 @@ SimRobotState* SimRobot::get_state() { } common::Pose SimRobot::get_cartesian_position() { + return this->get_cartesian_flange_position() * cfg.tcp_offset; +} + +common::Pose SimRobot::get_cartesian_flange_position() { Eigen::Matrix rotation( this->sim->d->site_xmat + 9 * this->ids.attachment_site); Eigen::Vector3d translation(this->sim->d->site_xpos + 3 * this->ids.attachment_site); common::Pose attachment_site(Eigen::Matrix3d(rotation), translation); - return this->to_pose_in_robot_coordinates(attachment_site) * cfg.tcp_offset; + return this->to_pose_in_robot_coordinates(attachment_site); } void SimRobot::set_joint_position(const common::VectorXd& q) { diff --git a/src/sim/SimRobot.h b/src/sim/SimRobot.h index 99b23184..5270bd67 100644 --- a/src/sim/SimRobot.h +++ b/src/sim/SimRobot.h @@ -66,6 +66,7 @@ class SimRobot : public common::Robot { SimRobotConfig* get_config() override; SimRobotState* get_state() override; common::Pose get_cartesian_position() override; + common::Pose get_cartesian_flange_position() override; void set_joint_position(const common::VectorXd& q) override; common::VectorXd get_joint_position() override; void move_home() override;