#[derive(
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
typed_builder::TypedBuilder,
)]
pub struct TaskConfig {
#[builder(default = String::from("sample_collection"))]
pub task: String,
#[builder(default = Workflow::Teleop)]
pub workflow: Workflow,
#[builder(default = Domain::Moon)]
pub domain: Domain,
#[builder(default = None)]
pub scenery: Option<String>,
#[builder(default = String::from("franka"))]
pub robot: String,
#[builder(default = None)]
pub end_effector: Option<String>,
#[builder(default = None)]
pub payload: Option<String>,
#[builder(default = 0)]
pub seed: u64,
#[builder(default = 1)]
pub num_envs: u64,
#[builder(default = false)]
pub stack_envs: bool,
#[builder(default = true)]
pub hide_ui: bool,
#[builder(default = vec![TeleopDevice::Keyboard, TeleopDevice::Spacemouse, TeleopDevice::Ros, TeleopDevice::Haptic])]
pub teleop_devices: Vec<TeleopDevice>,
#[builder(default = vec![InterfaceType::Gui])]
pub interfaces: Vec<InterfaceType>,
#[builder(default = vec![])]
pub extras: Vec<String>,
}
impl TaskConfig {
pub fn set_exec_env(&self, mut exec: subprocess::Exec) -> subprocess::Exec {
exec = exec.arg("agent");
match self.workflow {
Workflow::Teleop => {
exec = exec.arg(self.workflow.to_string());
exec = exec.arg("--teleop_device");
exec = exec.args(
&self
.teleop_devices
.iter()
.map(|device| device.to_string().to_lowercase())
.collect::<Vec<String>>(),
);
}
_ => {
exec = exec.arg(self.workflow.to_string());
}
}
if self.hide_ui {
exec = exec.arg("--hide_ui");
}
if !self.interfaces.is_empty() {
exec = exec.arg("--interface");
exec = exec.args(
&self
.interfaces
.iter()
.map(|interface| interface.to_string().to_lowercase())
.collect::<Vec<String>>(),
);
}
exec = exec.args(&["--task", &self.task]);
exec = exec.arg(format!(
"env.domain={}",
self.domain.to_string().to_lowercase()
));
exec = exec.arg(format!("env.seed={}", self.seed));
exec = exec.arg(format!("env.num_envs={}", self.num_envs.max(1)));
exec = exec.arg(format!(
"env.stack={}",
self.stack_envs.to_string().to_lowercase()
));
if !self.extras.is_empty() {
exec = exec.args(&self.extras);
}
if let Some(scenery) = &self.scenery {
exec = exec.arg(format!("env.scenery={scenery}"));
}
let robot_arg = self.build_robot_arg();
exec = exec.arg(format!("env.robot={robot_arg}"));
exec = exec.env(
"DISPLAY",
std::env::var("SRB_DISPLAY").unwrap_or(":0".to_string()),
);
exec = exec.env(
"ROS_DOMAIN_ID",
std::env::var("ROS_DOMAIN_ID").unwrap_or("0".to_string()),
);
exec = exec.env(
"RMW_IMPLEMENTATION",
std::env::var("RMW_IMPLEMENTATION").unwrap_or("rmw_cyclonedds_cpp".to_string()),
);
exec
}
fn build_robot_arg(&self) -> String {
let mut final_robot_arg = String::new();
let has_base = !self.robot.is_empty() && self.robot != "default";
let ee_opt = self.end_effector.as_deref();
let payload_opt = self.payload.as_deref();
let has_ee = ee_opt.is_some_and(|s| !s.is_empty() && s != "none" && s != "default");
let has_payload =
payload_opt.is_some_and(|s| !s.is_empty() && s != "none" && s != "default");
let keep_default_ee = ee_opt == Some("default");
let keep_default_payload = payload_opt == Some("default");
let no_ee = ee_opt == Some("none");
let no_payload = payload_opt == Some("none");
if has_base {
final_robot_arg.push_str(&self.robot);
if has_ee || keep_default_ee {
final_robot_arg.push('+');
if has_ee {
final_robot_arg.push_str(self.end_effector.as_ref().unwrap());
}
} else if has_payload || keep_default_payload {
final_robot_arg.push('+');
if has_payload {
final_robot_arg.push_str(self.payload.as_ref().unwrap());
}
} else if no_ee || no_payload {
}
} else if has_ee {
final_robot_arg.push('+');
final_robot_arg.push_str(self.end_effector.as_ref().unwrap());
} else if has_payload {
final_robot_arg.push('+');
final_robot_arg.push_str(self.payload.as_ref().unwrap());
}
if final_robot_arg.is_empty() || final_robot_arg == "+" {
String::new()
} else {
final_robot_arg
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum Domain {
Asteroid,
Earth,
Mars,
Moon,
Orbit,
}
impl std::str::FromStr for Domain {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"asteroid" => Ok(Self::Asteroid),
"earth" => Ok(Self::Earth),
"mars" => Ok(Self::Mars),
"moon" => Ok(Self::Moon),
"orbit" => Ok(Self::Orbit),
_ => Err(format!("Invalid Domain: {s}")),
}
}
}
impl Domain {
#[must_use]
pub fn gravity_magnitude(self) -> f64 {
match self {
Self::Asteroid => 0.14219,
Self::Earth => 9.80665,
Self::Mars => 3.72076,
Self::Moon => 1.62496,
Self::Orbit => 0.0,
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum RobotType {
Manipulator,
MobileManipulator,
MobileRobot,
}
impl std::str::FromStr for RobotType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"manipulator" => Ok(RobotType::Manipulator),
"mobile_manipulator" => Ok(RobotType::MobileManipulator),
"mobile_robot" => Ok(RobotType::MobileRobot),
_ => Err(format!("Invalid RobotType string: {s}")),
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum SceneryType {
Extravehicular,
Structure,
Subterrane,
Terrain,
}
impl std::str::FromStr for SceneryType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"extravehicular" => Ok(SceneryType::Extravehicular),
"structure" => Ok(SceneryType::Structure),
"subterrane" => Ok(SceneryType::Subterrane),
"terrain" => Ok(SceneryType::Terrain),
_ => Err(format!("Invalid SceneryType string: {s}")),
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum ObjectType {
Common,
Light,
Payload,
Pedestal,
Tool,
}
impl std::str::FromStr for ObjectType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"common" => Ok(ObjectType::Common),
"light" => Ok(ObjectType::Light),
"payload" => Ok(ObjectType::Payload),
"pedestal" => Ok(ObjectType::Pedestal),
"tool" => Ok(ObjectType::Tool),
_ => Err(format!("Invalid ObjectType: {s}")),
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum Workflow {
Rand,
Teleop,
Zero,
}
impl std::str::FromStr for Workflow {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"rand" => Ok(Self::Rand),
"teleop" => Ok(Self::Teleop),
"zero" => Ok(Self::Zero),
_ => Err(format!("Invalid Workflow: {s}")),
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum TeleopDevice {
Keyboard,
Ros,
Gamepad,
Spacemouse,
Haptic,
}
impl std::str::FromStr for TeleopDevice {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"keyboard" => Ok(TeleopDevice::Keyboard),
"ros" => Ok(TeleopDevice::Ros),
"gamepad" => Ok(TeleopDevice::Gamepad),
"spacemouse" => Ok(TeleopDevice::Spacemouse),
"haptic" => Ok(TeleopDevice::Haptic),
_ => Err(format!("Invalid TeleopDevice: {s}")),
}
}
}
#[derive(
Copy,
Clone,
Debug,
Eq,
PartialEq,
Hash,
serde::Deserialize,
serde::Serialize,
strum::Display,
strum::EnumIter,
)]
#[serde(rename_all = "snake_case")]
pub enum InterfaceType {
Gui,
Ros,
}
impl std::str::FromStr for InterfaceType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"gui" => Ok(InterfaceType::Gui),
"ros" => Ok(InterfaceType::Ros),
_ => Err(format!("Invalid InterfaceType: {s}")),
}
}
}