def test_simulator_ctor(self):
# Create simple system.
system = ConstantVectorSource([1])
def check_output(context):
# Check number of output ports and value for a given context.
output = system.AllocateOutput(context)
self.assertEquals(output.get_num_ports(), 1)
system.CalcOutput(context, output)
value = output.get_vector_data(0).get_value()
self.assertTrue(np.allclose([1], value))
# Create simulator with basic constructor.
simulator = Simulator(system)
simulator.Initialize()
simulator.set_target_realtime_rate(0)
simulator.set_publish_every_time_step(True)
self.assertTrue(
simulator.get_context() is simulator.get_mutable_context())
check_output(simulator.get_context())
simulator.StepTo(1)
# Create simulator specifying context.
context = system.CreateDefaultContext()
# @note `simulator` now owns `context`.
simulator = Simulator(system, context)
self.assertTrue(simulator.get_context() is context)
check_output(context)
simulator.StepTo(1)
def test_simulator_ctor(self):
# Create simple system.
system = ConstantVectorSource([1])
def check_output(context):
# Check number of output ports and value for a given context.
output = system.AllocateOutput(context)
self.assertEquals(output.get_num_ports(), 1)
system.CalcOutput(context, output)
value = output.get_vector_data(0).get_value()
self.assertTrue(np.allclose([1], value))
# Create simulator with basic constructor.
simulator = Simulator(system)
simulator.Initialize()
simulator.set_target_realtime_rate(0)
simulator.set_publish_every_time_step(True)
self.assertTrue(simulator.get_context() is
simulator.get_mutable_context())
check_output(simulator.get_context())
simulator.StepTo(1)
# Create simulator specifying context.
context = system.CreateDefaultContext()
# @note `simulator` now owns `context`.
simulator = Simulator(system, context)
self.assertTrue(simulator.get_context() is context)
check_output(context)
simulator.StepTo(1)
def experiment(start_state, force_schedule, sleeptime=None):
global experiment_count
experiment_count += 1
if experiment_count % 10 == 0:
print(f"{experiment_count} experiments run so far...")
builder = DiagramBuilder()
mbp, sg = make_mbp(builder)
mbp.Finalize()
DrakeVisualizer.AddToBuilder(builder, sg, lcm=global_lcm_ftw)
thrusters = add_thrusters(builder, mbp)
breaks = [0]
for _, t in force_schedule:
breaks.append(breaks[-1] + t)
forces = np.array([f for f, t in force_schedule] + [force_schedule[-1][0]])
force_traj = PiecewisePolynomial.ZeroOrderHold(breaks, forces.transpose())
controller = builder.AddSystem(TrajectorySource(force_traj))
builder.Connect(controller.get_output_port(),
thrusters.get_command_input_port())
diagram = builder.Build()
simulator = Simulator(diagram)
mbp_context = diagram.GetSubsystemContext(mbp, simulator.get_context())
mbp.SetPositionsAndVelocities(mbp_context, start_state)
for step in range(int(1000 * breaks[-1])):
simulator.AdvanceTo(0.001 * step)
mbp_context = diagram.GetSubsystemContext(mbp, simulator.get_context())
if sleeptime: time.sleep(sleeptime)
return mbp.GetPositionsAndVelocities(
diagram.GetSubsystemContext(mbp, simulator.get_context()))
def simulate_diagram(diagram, ball_paddle_plant, state_logger,
ball_init_position, ball_init_velocity, simulation_time,
target_realtime_rate):
q_init_val = np.array([
1, 0, 0, 0, ball_init_position[0], ball_init_position[1],
ball_init_position[2]
])
v_init_val = np.hstack((np.zeros(3), ball_init_velocity))
qv_init_val = np.concatenate((q_init_val, v_init_val))
simulator_config = SimulatorConfig(
target_realtime_rate=target_realtime_rate,
publish_every_time_step=True)
simulator = Simulator(diagram)
ApplySimulatorConfig(simulator_config, simulator)
plant_context = diagram.GetSubsystemContext(ball_paddle_plant,
simulator.get_context())
ball_paddle_plant.SetPositionsAndVelocities(plant_context, qv_init_val)
simulator.get_mutable_context().SetTime(0)
state_log = state_logger.FindMutableLog(simulator.get_mutable_context())
state_log.Clear()
simulator.Initialize()
simulator.AdvanceTo(boundary_time=simulation_time)
PrintSimulatorStatistics(simulator)
return state_log.sample_times(), state_log.data()
def test_subgraph_add_to_copying(self):
"""Ensures that index ordering is generally the same when copying a
plant using a MultibodyPlantSubgraph.add_to."""
# TODO(eric.cousineau): Increas number of bodies for viz, once
# `create_manual_map` can acommodate it.
plant_a, scene_graph_a, _ = self.make_arbitrary_multibody_stuff(
num_bodies=1)
# Check for general ordering with full subgraph "cloning".
builder_b = DiagramBuilder()
plant_b, scene_graph_b = AddMultibodyPlantSceneGraph(
builder_b, plant_a.time_step())
subgraph_a = mut.MultibodyPlantSubgraph(
mut.get_elements_from_plant(plant_a, scene_graph_a))
a_to_b = subgraph_a.add_to(plant_b, scene_graph_b)
plant_b.Finalize()
util.assert_plant_equals(plant_a, scene_graph_a, plant_b,
scene_graph_b)
a_to_b_expected = self.create_manual_map(plant_a, scene_graph_a,
plant_b, scene_graph_b)
self.assertEqual(a_to_b, a_to_b_expected)
if VISUALIZE:
for model in me.get_model_instances(plant_b):
util.build_with_no_control(builder_b, plant_b, model)
print("test_subgraph_add_to_copying")
DrakeVisualizer.AddToBuilder(builder_b, scene_graph_b)
diagram = builder_b.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.)
simulator.Initialize()
diagram.Publish(simulator.get_context())
simulator.AdvanceTo(1.)
def _process_event(self, dut):
# Use a Simulator to invoke the update event on `dut`. (Wouldn't it be
# nice if the Systems API was simple enough that we could apply events
# without calling a Simulator!)
simulator = Simulator(dut)
simulator.AdvanceTo(0.00025) # Arbitrary positive value.
return simulator.get_context().Clone()
def main():
# Simulate with doubles.
builder = DiagramBuilder()
source = builder.AddSystem(ConstantVectorSource([10.]))
adder = builder.AddSystem(SimpleAdder(100.))
builder.Connect(source.get_output_port(0), adder.get_input_port(0))
logger = builder.AddSystem(VectorLogSink(1))
builder.Connect(adder.get_output_port(0), logger.get_input_port(0))
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.AdvanceTo(1)
x = logger.FindLog(simulator.get_context()).data()
print("Output values: {}".format(x))
assert np.allclose(x, 110.)
# Compute outputs with AutoDiff and Symbolic.
for T in (AutoDiffXd, Expression):
adder_T = SimpleAdder_[T](100.)
context = adder_T.CreateDefaultContext()
adder_T.get_input_port().FixValue(context, [10.])
output = adder_T.AllocateOutput()
adder_T.CalcOutput(context, output)
# N.B. At present, you cannot get a reference to existing AutoDiffXd
# or Expression numpy arrays, so we will explictly copy the vector:
# https://github.com/RobotLocomotion/drake/issues/8116
value, = output.get_vector_data(0).CopyToVector()
assert isinstance(value, T)
print("Output from {}: {}".format(type(adder_T), repr(value)))
def main():
builder = DiagramBuilder()
source = builder.AddSystem(ConstantVectorSource([10.]))
logger = builder.AddSystem(VectorLogSink(1))
builder.Connect(source.get_output_port(0), logger.get_input_port(0))
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.AdvanceTo(1)
x = logger.FindLog(simulator.get_context()).data()
print("Output values: {}".format(x))
assert np.allclose(x, 10.)
def test_copying(self):
"""Ensures that index ordering is generally the same when copying a
plant using a subgraph."""
time_step = 0.01
builder_a = DiagramBuilder()
plant_a, scene_graph_a = AddMultibodyPlantSceneGraph(
builder_a, time_step)
util.add_arbitrary_multibody_stuff(plant_a)
# Ensure that this is "physically" valid.
plant_a.Finalize()
if VISUALIZE:
for model in me.get_model_instances(plant_a):
util.no_control(builder_a, plant_a, model)
print("test_composition_array_with_scene_graph")
ConnectDrakeVisualizer(builder_a, scene_graph_a)
diagram = builder_a.Build()
simulator = Simulator(diagram)
simulator.Initialize()
diagram.Publish(simulator.get_context())
simulator.AdvanceTo(1.)
# Checking for determinism, making a slight change to trigger an error.
for slight_difference in [False, True]:
builder_b = DiagramBuilder()
plant_b, scene_graph_b = AddMultibodyPlantSceneGraph(
builder_b, time_step)
util.add_arbitrary_multibody_stuff(
plant_b, slight_difference=slight_difference)
plant_b.Finalize()
if not slight_difference:
util.assert_plant_equals(plant_a, scene_graph_a, plant_b,
scene_graph_b)
else:
with self.assertRaises(AssertionError):
util.assert_plant_equals(plant_a, scene_graph_a, plant_b,
scene_graph_b)
# Check for general ordering with full subgraph "cloning".
builder_b = DiagramBuilder()
plant_b, scene_graph_b = AddMultibodyPlantSceneGraph(
builder_b, time_step)
subgraph_a = mut.MultibodyPlantSubgraph(
mut.get_elements_from_plant(plant_a, scene_graph_a))
subgraph_a.add_to(plant_b, scene_graph_b)
plant_b.Finalize()
util.assert_plant_equals(plant_a, scene_graph_a, plant_b,
scene_graph_b)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--target_realtime_rate",
type=float,
default=1.0,
help="Desired rate relative to real time. See documentation for "
"Simulator::set_target_realtime_rate() for details.")
parser.add_argument("--duration",
type=float,
default=np.inf,
help="Desired duration of the simulation in seconds.")
parser.add_argument(
"--hardware",
action='store_true',
help="Use the ManipulationStationHardwareInterface instead of an "
"in-process simulation.")
parser.add_argument("--test",
action='store_true',
help="Disable opening the gui window for testing.")
parser.add_argument(
'--setup',
type=str,
default='manipulation_class',
help="The manipulation station setup to simulate. ",
choices=['manipulation_class', 'clutter_clearing', 'planar'])
parser.add_argument(
"-w",
"--open-window",
dest="browser_new",
action="store_const",
const=1,
default=None,
help="Open the MeshCat display in a new browser window.")
args = parser.parse_args()
builder = DiagramBuilder()
# NOTE: the meshcat instance is always created in order to create the
# teleop controls (joint sliders and open/close gripper button). When
# args.hardware is True, the meshcat server will *not* display robot
# geometry, but it will contain the joint sliders and open/close gripper
# button in the "Open Controls" tab in the top-right of the viewing server.
meshcat = Meshcat()
if args.hardware:
# TODO(russt): Replace this hard-coded camera serial number with a
# config file.
camera_ids = ["805212060544"]
station = builder.AddSystem(
ManipulationStationHardwareInterface(camera_ids))
station.Connect(wait_for_cameras=False)
else:
station = builder.AddSystem(ManipulationStation())
# Initializes the chosen station type.
if args.setup == 'manipulation_class':
station.SetupManipulationClassStation()
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/" +
"061_foam_brick.sdf",
RigidTransform(RotationMatrix.Identity(), [0.6, 0, 0]))
elif args.setup == 'clutter_clearing':
station.SetupClutterClearingStation()
ycb_objects = CreateClutterClearingYcbObjectList()
for model_file, X_WObject in ycb_objects:
station.AddManipulandFromFile(model_file, X_WObject)
elif args.setup == 'planar':
station.SetupPlanarIiwaStation()
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/" +
"061_foam_brick.sdf",
RigidTransform(RotationMatrix.Identity(), [0.6, 0, 0]))
station.Finalize()
geometry_query_port = station.GetOutputPort("geometry_query")
DrakeVisualizer.AddToBuilder(builder, geometry_query_port)
meshcat_visualizer = MeshcatVisualizerCpp.AddToBuilder(
builder=builder,
query_object_port=geometry_query_port,
meshcat=meshcat)
if args.setup == 'planar':
meshcat.Set2dRenderMode()
pyplot_visualizer = ConnectPlanarSceneGraphVisualizer(
builder, station.get_scene_graph(), geometry_query_port)
if args.browser_new is not None:
url = meshcat.web_url()
webbrowser.open(url=url, new=args.browser_new)
teleop = builder.AddSystem(
JointSliders(meshcat=meshcat, plant=station.get_controller_plant()))
num_iiwa_joints = station.num_iiwa_joints()
filter = builder.AddSystem(
FirstOrderLowPassFilter(time_constant=2.0, size=num_iiwa_joints))
builder.Connect(teleop.get_output_port(0), filter.get_input_port(0))
builder.Connect(filter.get_output_port(0),
station.GetInputPort("iiwa_position"))
wsg_buttons = builder.AddSystem(SchunkWsgButtons(meshcat=meshcat))
builder.Connect(wsg_buttons.GetOutputPort("position"),
station.GetInputPort("wsg_position"))
builder.Connect(wsg_buttons.GetOutputPort("force_limit"),
station.GetInputPort("wsg_force_limit"))
# When in regression test mode, log our joint velocities to later check
# that they were sufficiently quiet.
if args.test:
iiwa_velocities = builder.AddSystem(VectorLogSink(num_iiwa_joints))
builder.Connect(station.GetOutputPort("iiwa_velocity_estimated"),
iiwa_velocities.get_input_port(0))
else:
iiwa_velocities = None
diagram = builder.Build()
simulator = Simulator(diagram)
# This is important to avoid duplicate publishes to the hardware interface:
simulator.set_publish_every_time_step(False)
station_context = diagram.GetMutableSubsystemContext(
station, simulator.get_mutable_context())
station.GetInputPort("iiwa_feedforward_torque").FixValue(
station_context, np.zeros(num_iiwa_joints))
# If the diagram is only the hardware interface, then we must advance it a
# little bit so that first LCM messages get processed. A simulated plant is
# already publishing correct positions even without advancing, and indeed
# we must not advance a simulated plant until the sliders and filters have
# been initialized to match the plant.
if args.hardware:
simulator.AdvanceTo(1e-6)
# Eval the output port once to read the initial positions of the IIWA.
q0 = station.GetOutputPort("iiwa_position_measured").Eval(station_context)
teleop.SetPositions(q0)
filter.set_initial_output_value(
diagram.GetMutableSubsystemContext(filter,
simulator.get_mutable_context()),
q0)
simulator.set_target_realtime_rate(args.target_realtime_rate)
simulator.AdvanceTo(args.duration)
# Ensure that our initialization logic was correct, by inspecting our
# logged joint velocities.
if args.test:
iiwa_velocities_log = iiwa_velocities.FindLog(simulator.get_context())
for time, qdot in zip(iiwa_velocities_log.sample_times(),
iiwa_velocities_log.data().transpose()):
# TODO(jwnimmer-tri) We should be able to do better than a 40
# rad/sec limit, but that's the best we can enforce for now.
if qdot.max() > 0.1:
print(f"ERROR: large qdot {qdot} at time {time}")
sys.exit(1)
class PusherSliderEnv(gym.Env):
metadata = {"render.modes": []}
def __init__(self, config=None):
if config is None:
config_path = os.path.join(os.path.dirname(__file__),
"config.yaml")
config = yaml.safe_load(open(config_path, 'r'))
self._config = config
self._step_dt = config["step_dt"]
self._model_name = config["model_name"]
self._sim_diagram = DrakeSimDiagram(config)
mbp = self._sim_diagram.mbp
builder = self._sim_diagram.builder
# === Add table =====
dims = config["table"]["size"]
color = np.array(config["table"]["color"])
friction_params = config["table"]["coulomb_friction"]
box_shape = Box(*dims)
X_W_T = RigidTransform(p=np.array([0., 0., -dims[2] / 2.]))
# This rigid body will be added to the world model instance since
# the model instance is not specified.
box_body = mbp.AddRigidBody(
"table",
SpatialInertia(mass=1.0,
p_PScm_E=np.array([0., 0., 0.]),
G_SP_E=UnitInertia(1.0, 1.0, 1.0)))
mbp.WeldFrames(mbp.world_frame(), box_body.body_frame(), X_W_T)
mbp.RegisterVisualGeometry(box_body, RigidTransform(), box_shape,
"table_vis", color)
mbp.RegisterCollisionGeometry(box_body, RigidTransform(), box_shape,
"table_collision",
CoulombFriction(*friction_params))
# === Add pusher ====
parser = Parser(mbp, self._sim_diagram.sg)
self._pusher_model_id = parser.AddModelFromFile(
path_util.pusher_peg, "pusher")
base_link = mbp.GetBodyByName("base", self._pusher_model_id)
mbp.WeldFrames(mbp.world_frame(), base_link.body_frame())
def pusher_port_func():
actuation_input_port = mbp.get_actuation_input_port(
self._pusher_model_id)
state_output_port = mbp.get_state_output_port(
self._pusher_model_id)
builder.ExportInput(actuation_input_port, "pusher_actuation_input")
builder.ExportOutput(state_output_port, "pusher_state_output")
self._sim_diagram.finalize_functions.append(pusher_port_func)
# === Add slider ====
parser = Parser(mbp, self._sim_diagram.sg)
self._slider_model_id = parser.AddModelFromFile(
path_util.model_paths[self._model_name], self._model_name)
def slider_port_func():
state_output_port = mbp.get_state_output_port(
self._slider_model_id)
builder.ExportOutput(state_output_port, "slider_state_output")
self._sim_diagram.finalize_functions.append(slider_port_func)
if "rgbd_sensors" in config and config["rgbd_sensors"]["enabled"]:
self._sim_diagram.add_rgbd_sensors_from_config(config)
if "visualization" in config:
if not config["visualization"]:
pass
elif config["visualization"] == "meshcat":
self._sim_diagram.connect_to_meshcat()
elif config["visualization"] == "drake_viz":
self._sim_diagram.connect_to_drake_visualizer()
else:
raise ValueError("Unknown visualization:",
config["visualization"])
self._sim_diagram.finalize()
builder = DiagramBuilder()
builder.AddSystem(self._sim_diagram)
pid = builder.AddSystem(
PidController(kp=[0, 0], kd=[1000, 1000], ki=[0, 0]))
builder.Connect(self._sim_diagram.GetOutputPort("pusher_state_output"),
pid.get_input_port_estimated_state())
builder.Connect(
pid.get_output_port_control(),
self._sim_diagram.GetInputPort("pusher_actuation_input"))
builder.ExportInput(pid.get_input_port_desired_state(),
"pid_input_port_desired_state")
self._diagram = builder.Build()
self._pid_desired_state_port = self._diagram.get_input_port(0)
self._simulator = Simulator(self._diagram)
self.reset()
self.action_space = spaces.Box(low=-1.,
high=1.,
shape=(2, ),
dtype=np.float32)
# TODO: Observation space for images is currently too loose
self.observation_space = convert_observation_to_space(
self.get_observation())
def get_observation(self, context=None):
assert self._sim_diagram.is_finalized()
if context is None:
context = self._simulator.get_context()
context = self._diagram.GetMutableSubsystemContext(
self._sim_diagram, context)
mbp_context = self._sim_diagram.GetMutableSubsystemContext(
self._sim_diagram.mbp, context)
obs = dict()
obs["time"] = context.get_time()
obs["pusher"] = dict()
obs["pusher"]["position"] = np.copy(
self._sim_diagram.mbp.GetPositions(mbp_context,
self._pusher_model_id))
obs["pusher"]["velocity"] = np.copy(
self._sim_diagram.mbp.GetVelocities(mbp_context,
self._pusher_model_id))
q = np.copy(
self._sim_diagram.mbp.GetPositions(mbp_context,
self._slider_model_id))
v = np.copy(
self._sim_diagram.mbp.GetVelocities(mbp_context,
self._slider_model_id))
obs["slider"] = dict()
obs["slider"]["position"] = {
'translation': q[4:],
'quaternion': q[0:4],
'raw': q
}
obs["slider"]["velocity"] = {
'linear': v[3:],
'angular': v[0:3],
'raw': v
}
if self._config["rgbd_sensors"]["enabled"]:
obs["images"] = self._sim_diagram.get_image_observations(context)
return obs
def pusher_slider_on_table(self, context=None):
if context is None:
context = self._simulator.get_context()
context = self._diagram.GetMutableSubsystemContext(
self._sim_diagram, context)
mbp_context = self._sim_diagram.GetMutableSubsystemContext(
self._sim_diagram.mbp, context)
table_dim = np.array(self._config["table"]["size"])
q_pusher = self._sim_diagram.mbp.GetPositions(mbp_context,
self._pusher_model_id)
q_slider = self._sim_diagram.mbp.GetPositions(
mbp_context, self._slider_model_id)[4:6]
tol = 0.03
high_edge = table_dim[:2] / 2.0 - tol
low_edge = -high_edge
return ((low_edge < q_pusher) & (q_pusher < high_edge) &
(low_edge < q_slider) & (q_slider < high_edge)).all()
def step(self, action, dt=None):
assert self._sim_diagram.is_finalized()
assert len(action) == 2
if dt is None:
dt = self._step_dt
# the time to advance to
t_advance = self._simulator.get_context().get_time() + dt
context = self._simulator.get_mutable_context()
# set the value of the PID controller input port
pid_setpoint = np.concatenate((np.zeros(2), action))
self._pid_desired_state_port.FixValue(context, pid_setpoint)
# simulate and take observation
self._simulator.AdvanceTo(t_advance)
obs = self.get_observation()
reward = 0.
done = not self.pusher_slider_on_table()
info = {}
return obs, reward, done, info
# TODO: Make reset position non-deterministic?
def reset(self, q_pusher=None, q_slider=None):
assert self._sim_diagram.is_finalized()
if q_pusher is None:
q_pusher = np.array([-0.1, 0.])
if q_slider is None:
q_slider = np.array([1., 0., 0., 0., 0., 0., 0.05])
pid_setpoint = np.zeros(4)
context = self._simulator.get_mutable_context()
mbp_context = self._diagram.GetMutableSubsystemContext(
self._sim_diagram.mbp, context)
context.SetTime(0.)
self._sim_diagram.mbp.SetPositions(mbp_context, self._pusher_model_id,
q_pusher)
self._sim_diagram.mbp.SetVelocities(mbp_context, self._pusher_model_id,
np.zeros(2))
self._sim_diagram.mbp.SetPositions(mbp_context, self._slider_model_id,
q_slider)
self._sim_diagram.mbp.SetVelocities(mbp_context, self._slider_model_id,
np.zeros(6))
self._pid_desired_state_port.FixValue(context, pid_setpoint)
self._simulator.Initialize()
return self.get_observation()
#TODO: Implement as necessary
def render(self, mode="human"):
pass
def close(self):
pass
#TODO: Implement?
def seed(self, seed=None):
pass
def get_labels_to_mask(self):
"""
Returns a set of labels that represent slider object in the scene
:return: TinyDB database
"""
# write sql query
label_db = self._sim_diagram.get_label_db()
# return all objects that have matching model_instance_id in the list
Label = Query()
labels_to_mask = label_db.search(
Label.model_instance_id.one_of(int(self._slider_model_id)))
# create new TinyDB with just those
mask_db = TinyDB(storage=MemoryStorage)
mask_db.insert_multiple(labels_to_mask)
return mask_db
class DrakeMugsEnv(Env):
def __init__(self, config, visualize=True):
Env.__init__(self)
self._config = config
self._step_dt = config['env']['step_dt']
self._model_name = "mug"
# setup DPS wrapper
self._diagram_wrapper = DrakePusherSliderDiagramWrapper(config)
# setup the simulator
# add procedurally generated table
env_utils.add_procedurally_generated_table(self.diagram_wrapper.mbp,
config['env']['table']),
self.diagram_wrapper.add_pusher()
self.add_object_model()
self.diagram_wrapper.finalize()
self.diagram_wrapper.export_ports()
if visualize:
self.diagram_wrapper.connect_to_meshcat()
self.diagram_wrapper.connect_to_drake_visualizer()
self.diagram_wrapper.add_sensors_from_config(config)
self.diagram_wrapper.build()
# records port indices
self._port_idx = dict()
# add controller and other stuff
builder = DiagramBuilder()
self._builder = builder
# print("type(self.diagram_wrapper.diagram)", type(self.diagram_wrapper.diagram))
builder.AddSystem(self.diagram_wrapper.diagram)
# need to connect actuator ports
# set the controller gains
pid_data = self.diagram_wrapper.add_pid_controller(builder=builder)
self._port_idx["pid_input_port_desired_state"] = pid_data[
'pid_input_port_index']
diagram = builder.Build()
self._diagram = diagram
self._pid_input_port_desired_state = self._diagram.get_input_port(
self._port_idx["pid_input_port_desired_state"])
# setup simulator
context = diagram.CreateDefaultContext()
self._simulator = Simulator(self._diagram, context)
self._sim_initialized = False
self._context = context
# reset env
self.reset()
@property
def diagram_wrapper(self):
return self._diagram_wrapper
def add_object_model(self):
# add mug
# sdf_path = "anzu_mugs/big_mug-corelle_mug-6.sdf"
# sdf_path = "anzu_mugs/big_mug-small_mug-0.sdf"
# sdf_path = "anzu_mugs/corelle_mug-small_mug-8.sdf"
# sdf_path = "manual_babson_11oz_mug/manual_babson_11oz_mug.sdf"
# sdf_path = os.path.join(LARGE_SIM_ASSETS_ROOT, sdf_path)
#
self._object_name = "mug"
sdf_file_fullpath = os.path.join(get_data_root(),
self.config['env']['model']['sdf'])
model_color = self.config['env']['model']['color']
# output_dir = "/home/manuelli/data/key_dynam/sandbox/sdf_helper"
output_dir = None
sdf_data = SDFHelper.create_sdf_specific_color(
sdf_file_fullpath=sdf_file_fullpath,
color=model_color,
output_dir=output_dir)
self.diagram_wrapper.add_model_from_sdf(self._object_name,
sdf_data['sdf_file'])
def reset(self):
"""
Sets up the environment according to the config
:return:
:rtype:
"""
context = self.get_mutable_context()
# self.simulator.SetDefaultContext(context)
# self.diagram.SetDefaultContext(context)
# # put the ycb model slightly above the table
q_object = np.array([1, 0, 0, 0, 0, 0, 0.1])
self.set_object_position(context, q_object)
# self.set_slider_position(context, q_ycb_model)
# self.set_slider_velocity(context, np.zeros(6))
# set pusher initial position
q_pusher = np.array([-0.1, 0.])
self.set_pusher_position(context, q_pusher)
self.set_pusher_velocity(context, np.zeros(2))
# set PID setpoint to zero
pid_setpoint = np.zeros(4)
self._pid_input_port_desired_state.FixValue(context, pid_setpoint)
self._sim_initialized = False
def step(self, action, dt=None):
assert action.size == 2
if not self._sim_initialized:
self._simulator.Initialize()
self._sim_initialized = True
if dt is None:
dt = self._step_dt
# the time to advance to
t_advance = self.get_context().get_time() + dt
context = self.get_mutable_context()
# set the value of the PID controller input port
pid_setpoint = np.concatenate((np.zeros(2), action))
self._pid_input_port_desired_state.FixValue(context, pid_setpoint)
# self._velocity_controller.set_velocity(action)
self._simulator.AdvanceTo(t_advance)
# set the sim time from the above context to avoid floating point errors
obs = self.get_observation()
# just placeholders for now
reward = 0.
done = False
info = None
return obs, reward, done, info
def render(self, mode='rgb'):
pass
def get_observation(
self,
context=None, # the specific context to use, otherwise none
):
"""
Gets obesrvations (optionally including images)
:param context:
:type context:
:return:
:rtype:
"""
if context is None:
context = self.get_context()
obs = dict()
obs['slider'] = self.diagram_wrapper.get_model_state(
self.diagram, context, self._object_name)
obs['object'] = obs['slider']
obs['pusher'] = self.diagram_wrapper.get_pusher_state(
self.diagram, context)
obs['sim_time'] = context.get_time()
if self._config["env"]["observation"]["rgbd"]:
obs["images"] = self.diagram_wrapper.get_image_observations(
self.diagram, context)
return obs
def get_mutable_context(self):
return self._simulator.get_mutable_context()
def get_context(self):
return self._simulator.get_context()
def get_rgb_image(
self,
sensor_name: str,
context,
):
sensor = self.diagram_wrapper._rgbd_sensors[sensor_name]
print("type(sensor)", type(sensor))
sensor_context = self.diagram.GetSubsystemContext(sensor, context)
rgb_img_PIL = drake_image_utils.get_color_image(sensor, sensor_context)
return rgb_img_PIL
def set_pusher_position(self, context, q):
"""
Sets the pusher position
:param context:
:type context:
:return:
:rtype:
"""
assert q.size == 2
mbp = self.diagram_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetPositions(mbp_context, self.diagram_wrapper.models['pusher'], q)
def set_pusher_velocity(self, context, v):
assert v.size == 2
mbp = self.diagram_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetVelocities(mbp_context, self.diagram_wrapper.models['pusher'],
v)
def set_object_position(self, context, q):
assert q.size == 7
# check that first 4 are a quaternion
if abs(np.linalg.norm(q[0:4]) - 1) > 1e-2:
print("q[0:4]", q[0:4])
print("np.linalg.norm(q[0:4])", np.linalg.norm(q[0:4]))
assert (abs(np.linalg.norm(q[0:4]) - 1) <
1e-2), "q[0:4] is not a valid quaternion"
mbp = self.diagram_wrapper.mbp
env_utils.set_model_position(diagram=self.diagram,
context=context,
mbp=mbp,
model_name=self._object_name,
q=q)
def set_object_velocity(self, context, v):
assert v.size == 6
mbp = self.diagram_wrapper.mbp
env_utils.set_model_velocity(diagram=self.diagram,
context=context,
mbp=mbp,
model_name=self._object_name,
v=v)
def get_zero_action(self):
"""
Returns a zero action
:return:
:rtype:
"""
return np.zeros(2)
@property
def simulator(self):
return self._simulator
@property
def diagram(self):
return self._diagram
def camera_pose(self, camera_name): # [4,4] homogeneous transform
pose_dict = self.config['env']['rgbd_sensors']['sensor_list'][
camera_name]
return transform_utils.transform_from_pose_dict(pose_dict)
@property
def config(self):
return self._config
def get_metadata(self):
"""
Returns a bunch of data that is useful to be logged
:return: dict()
:rtype:
"""
label_db = self.diagram_wrapper.get_label_db()
mask_db = self.diagram_wrapper.get_labels_to_mask()
# note saving TinyDB object with pickle causes error on load
# this is a workaround, can easily re-create TinyDB object
# when we load this
return {'label_db': label_db.all(), 'mask_db': mask_db.all()}
def get_labels_to_mask_list(self):
"""
Returns a list of labels that should be masked
:return:
:rtype:
"""
return self.diagram_wrapper.get_labels_to_mask_list()
def pusher_within_boundary(self, context=None):
"""
Return true if pusher within boundary of table + tolerance
:return:
:rtype:
"""
if context is None:
context = self.get_context()
tol = 0.03 # 3 cm
pusher_state = self.diagram_wrapper.get_pusher_state(
self.diagram, context)
pos = pusher_state['position'] # x,y
dims = np.array(self.config['env']['table']['size'])
high = dims[0:2] / 2.0 - tol
low = -high
in_range = ((low < pos) & (pos < high)).all()
return in_range
def slider_within_boundary(self, context=None):
"""
Return true if slider within boundary of table + tolerance
:return:
:rtype:
"""
if context is None:
context = self.get_context()
tol = 0.03 # 3 cm
slider_state = self.diagram_wrapper.get_model_state(
self.diagram, context, self._model_name)
pos = slider_state['position']['translation']
dims = np.array(self.config['env']['table']['size'])
high = np.zeros(3)
low = np.zeros(3)
high[0:2] = dims[0:2] / 2.0 - tol
high[2] = 100 # just large number so that it's nonbinding
low[0:2] = -high[0:2]
low[2] = 0
in_range = ((low < pos) & (pos < high)).all()
return in_range
def camera_K_matrix(self, camera_name): # np.array [4,4]
return DrakeMugsEnv.camera_K_matrix_from_config(
self.config, camera_name)
@staticmethod
def camera_K_matrix_from_config(config, camera_name): # np.array [4,4]
sensor_dict = config['env']['rgbd_sensors']['sensor_list'][camera_name]
width = sensor_dict['width']
height = sensor_dict['height']
fov_y = sensor_dict['fov_y']
return transform_utils.camera_K_matrix_from_FOV(width, height, fov_y)
def set_simulator_state_from_observation_dict(
self,
context,
d, # dictionary returned by DrakeSimEpisodeReader
):
# set slider position
q_slider = d['slider']['position']['raw']
v_slider = d['slider']['velocity']['raw']
self.set_object_position(context, q_slider)
self.set_object_velocity(context, v_slider)
# set pusher position
q_pusher = d['pusher']['position']
v_pusher = d['pusher']['velocity']
self.set_pusher_position(context, q_pusher)
self.set_pusher_velocity(context, v_pusher)
def set_initial_condition_from_dict(self, initial_cond):
context = self.get_mutable_context()
self.reset() # sets velocities to zero
self.set_pusher_position(context, initial_cond['q_pusher'])
self.set_object_position(context, initial_cond['q_object'])
self.step(np.array([0, 0]), dt=10.0) # allow box to drop down
def state_is_valid(self): # checks if state is valid
return (self.slider_within_boundary()
and self.pusher_within_boundary())
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--target_realtime_rate", type=float, default=1.0,
help="Desired rate relative to real time. See documentation for "
"Simulator::set_target_realtime_rate() for details.")
parser.add_argument(
"--duration", type=float, default=np.inf,
help="Desired duration of the simulation in seconds.")
parser.add_argument(
"--hardware", action='store_true',
help="Use the ManipulationStationHardwareInterface instead of an "
"in-process simulation.")
parser.add_argument(
"--test", action='store_true',
help="Disable opening the gui window for testing.")
parser.add_argument(
"--filter_time_const", type=float, default=0.1,
help="Time constant for the first order low pass filter applied to"
"the teleop commands")
parser.add_argument(
"--velocity_limit_factor", type=float, default=1.0,
help="This value, typically between 0 and 1, further limits the "
"iiwa14 joint velocities. It multiplies each of the seven "
"pre-defined joint velocity limits. "
"Note: The pre-defined velocity limits are specified by "
"iiwa14_velocity_limits, found in this python file.")
parser.add_argument(
'--setup', type=str, default='manipulation_class',
help="The manipulation station setup to simulate. ",
choices=['manipulation_class', 'clutter_clearing', 'planar'])
parser.add_argument(
'--schunk_collision_model', type=str, default='box',
help="The Schunk collision model to use for simulation. ",
choices=['box', 'box_plus_fingertip_spheres'])
parser.add_argument(
"-w", "--open-window", dest="browser_new",
action="store_const", const=1, default=None,
help=(
"Open the MeshCat display in a new browser window. NOTE: the "
"slider controls are available in the meshcat viewer by clicking "
"'Open Controls' in the top-right corner."))
args = parser.parse_args()
builder = DiagramBuilder()
# NOTE: the meshcat instance is always created in order to create the
# teleop controls (orientation sliders and open/close gripper button). When
# args.hardware is True, the meshcat server will *not* display robot
# geometry, but it will contain the joint sliders and open/close gripper
# button in the "Open Controls" tab in the top-right of the viewing server.
meshcat = Meshcat()
if args.hardware:
station = builder.AddSystem(ManipulationStationHardwareInterface())
station.Connect(wait_for_cameras=False)
else:
station = builder.AddSystem(ManipulationStation())
if args.schunk_collision_model == "box":
schunk_model = SchunkCollisionModel.kBox
elif args.schunk_collision_model == "box_plus_fingertip_spheres":
schunk_model = SchunkCollisionModel.kBoxPlusFingertipSpheres
# Initializes the chosen station type.
if args.setup == 'manipulation_class':
station.SetupManipulationClassStation(
schunk_model=schunk_model)
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/"
+ "061_foam_brick.sdf",
RigidTransform(RotationMatrix.Identity(), [0.6, 0, 0]))
elif args.setup == 'clutter_clearing':
station.SetupClutterClearingStation(
schunk_model=schunk_model)
ycb_objects = CreateClutterClearingYcbObjectList()
for model_file, X_WObject in ycb_objects:
station.AddManipulandFromFile(model_file, X_WObject)
elif args.setup == 'planar':
station.SetupPlanarIiwaStation(
schunk_model=schunk_model)
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/"
+ "061_foam_brick.sdf",
RigidTransform(RotationMatrix.Identity(), [0.6, 0, 0]))
station.Finalize()
# If using meshcat, don't render the cameras, since RgbdCamera
# rendering only works with drake-visualizer. Without this check,
# running this code in a docker container produces libGL errors.
geometry_query_port = station.GetOutputPort("geometry_query")
# Connect the meshcat visualizer.
meshcat_visualizer = MeshcatVisualizer.AddToBuilder(
builder=builder,
query_object_port=geometry_query_port,
meshcat=meshcat)
# Configure the planar visualization.
if args.setup == 'planar':
meshcat.Set2dRenderMode()
# Connect and publish to drake visualizer.
DrakeVisualizer.AddToBuilder(builder, geometry_query_port)
image_to_lcm_image_array = builder.AddSystem(
ImageToLcmImageArrayT())
image_to_lcm_image_array.set_name("converter")
for name in station.get_camera_names():
cam_port = (
image_to_lcm_image_array
.DeclareImageInputPort[PixelType.kRgba8U](
"camera_" + name))
builder.Connect(
station.GetOutputPort("camera_" + name + "_rgb_image"),
cam_port)
image_array_lcm_publisher = builder.AddSystem(
LcmPublisherSystem.Make(
channel="DRAKE_RGBD_CAMERA_IMAGES",
lcm_type=lcmt_image_array,
lcm=None,
publish_period=0.1,
use_cpp_serializer=True))
image_array_lcm_publisher.set_name("rgbd_publisher")
builder.Connect(
image_to_lcm_image_array.image_array_t_msg_output_port(),
image_array_lcm_publisher.get_input_port(0))
if args.browser_new is not None:
url = meshcat.web_url()
webbrowser.open(url=url, new=args.browser_new)
robot = station.get_controller_plant()
params = DifferentialInverseKinematicsParameters(robot.num_positions(),
robot.num_velocities())
time_step = 0.005
params.set_timestep(time_step)
# True velocity limits for the IIWA14 (in rad, rounded down to the first
# decimal)
iiwa14_velocity_limits = np.array([1.4, 1.4, 1.7, 1.3, 2.2, 2.3, 2.3])
if args.setup == 'planar':
# Extract the 3 joints that are not welded in the planar version.
iiwa14_velocity_limits = iiwa14_velocity_limits[1:6:2]
# The below constant is in body frame.
params.set_end_effector_velocity_gain([1, 0, 0, 0, 1, 1])
# Stay within a small fraction of those limits for this teleop demo.
factor = args.velocity_limit_factor
params.set_joint_velocity_limits((-factor*iiwa14_velocity_limits,
factor*iiwa14_velocity_limits))
differential_ik = builder.AddSystem(DifferentialIK(
robot, robot.GetFrameByName("iiwa_link_7"), params, time_step))
builder.Connect(differential_ik.GetOutputPort("joint_position_desired"),
station.GetInputPort("iiwa_position"))
teleop = builder.AddSystem(EndEffectorTeleop(
meshcat, args.setup == 'planar'))
filter = builder.AddSystem(
FirstOrderLowPassFilter(time_constant=args.filter_time_const, size=6))
builder.Connect(teleop.get_output_port(0), filter.get_input_port(0))
builder.Connect(filter.get_output_port(0),
differential_ik.GetInputPort("rpy_xyz_desired"))
wsg_buttons = builder.AddSystem(SchunkWsgButtons(meshcat=meshcat))
builder.Connect(wsg_buttons.GetOutputPort("position"),
station.GetInputPort("wsg_position"))
builder.Connect(wsg_buttons.GetOutputPort("force_limit"),
station.GetInputPort("wsg_force_limit"))
# When in regression test mode, log our joint velocities to later check
# that they were sufficiently quiet.
num_iiwa_joints = station.num_iiwa_joints()
if args.test:
iiwa_velocities = builder.AddSystem(VectorLogSink(num_iiwa_joints))
builder.Connect(station.GetOutputPort("iiwa_velocity_estimated"),
iiwa_velocities.get_input_port(0))
else:
iiwa_velocities = None
diagram = builder.Build()
simulator = Simulator(diagram)
# This is important to avoid duplicate publishes to the hardware interface:
simulator.set_publish_every_time_step(False)
station_context = diagram.GetMutableSubsystemContext(
station, simulator.get_mutable_context())
station.GetInputPort("iiwa_feedforward_torque").FixValue(
station_context, np.zeros(num_iiwa_joints))
# If the diagram is only the hardware interface, then we must advance it a
# little bit so that first LCM messages get processed. A simulated plant is
# already publishing correct positions even without advancing, and indeed
# we must not advance a simulated plant until the sliders and filters have
# been initialized to match the plant.
if args.hardware:
simulator.AdvanceTo(1e-6)
q0 = station.GetOutputPort("iiwa_position_measured").Eval(
station_context)
differential_ik.parameters.set_nominal_joint_position(q0)
teleop.SetPose(differential_ik.ForwardKinematics(q0))
filter.set_initial_output_value(
diagram.GetMutableSubsystemContext(
filter, simulator.get_mutable_context()),
teleop.get_output_port(0).Eval(diagram.GetMutableSubsystemContext(
teleop, simulator.get_mutable_context())))
differential_ik.SetPositions(diagram.GetMutableSubsystemContext(
differential_ik, simulator.get_mutable_context()), q0)
simulator.set_target_realtime_rate(args.target_realtime_rate)
simulator.AdvanceTo(args.duration)
# Ensure that our initialization logic was correct, by inspecting our
# logged joint velocities.
if args.test:
iiwa_velocities_log = iiwa_velocities.FindLog(simulator.get_context())
for time, qdot in zip(iiwa_velocities_log.sample_times(),
iiwa_velocities_log.data().transpose()):
# TODO(jwnimmer-tri) We should be able to do better than a 40
# rad/sec limit, but that's the best we can enforce for now.
if qdot.max() > 0.1:
print(f"ERROR: large qdot {qdot} at time {time}")
sys.exit(1)
class DrakePusherSliderEnv(Env):
def __init__(self, config, visualize=True):
Env.__init__(self)
self._config = config
self._step_dt = config['env']['step_dt']
self._model_name = config['env']['model_name']
# setup DPS wrapper
self._dps_wrapper = DrakePusherSliderDiagramWrapper(config)
# setup the simulator
self._dps_wrapper.add_procedurally_generated_table()
self._dps_wrapper.add_ycb_model_from_sdf(config['env']['model_name'])
self._dps_wrapper.add_pusher()
self._dps_wrapper.finalize()
self._dps_wrapper.export_ports()
if visualize:
self._dps_wrapper.connect_to_meshcat()
self._dps_wrapper.connect_to_drake_visualizer()
self._dps_wrapper.add_sensors_from_config(config)
self._dps_wrapper.build()
# records port indices
self._port_idx = dict()
# add controller and other stuff
builder = DiagramBuilder()
self._builder = builder
# print("type(self._dps_wrapper.diagram)", type(self._dps_wrapper.diagram))
builder.AddSystem(self._dps_wrapper.diagram)
# need to connect actuator ports
# set the controller gains
self.add_pid_controller()
diagram = builder.Build()
self._diagram = diagram
self._pid_input_port_desired_state = self.diagram.get_input_port(
self._port_idx["pid_input_port_desired_state"])
# setup simulator
context = diagram.CreateDefaultContext()
self._simulator = Simulator(diagram, context)
self._sim_initialized = False
self._context = context
# reset env
self.reset()
def add_pid_controller(self):
"""
Adds PID controller to system
:return:
:rtype:
"""
mbp = self._dps_wrapper.mbp
builder = self._builder
kp = [0, 0]
kd_gain = 1000
kd = [kd_gain] * 2
ki = [0, 0]
pid = builder.AddSystem(PidController(kp=kp, kd=kd, ki=ki))
# pusher state output port
pusher_model_idx = self._dps_wrapper.models['pusher']
pusher_state_output_port = mbp.get_state_output_port(pusher_model_idx)
pusher_actuation_input_port = mbp.get_actuation_input_port(
pusher_model_idx)
port_name = self._dps_wrapper.port_names["pusher_state_output"]
pusher_state_output_port = self._dps_wrapper.diagram.GetOutputPort(
port_name)
port_name = self._dps_wrapper.port_names["pusher_actuation_input"]
pusher_actuation_input_port = self._dps_wrapper.diagram.GetInputPort(
port_name)
builder.Connect(pusher_state_output_port,
pid.get_input_port_estimated_state())
builder.Connect(pid.get_output_port_control(),
pusher_actuation_input_port)
# make this visible externally
# will need to be connected up to SliderVelocityController(LeafSystem)
port_name = "pid_input_port_desired_state"
self._port_idx[port_name] = builder.ExportInput(
pid.get_input_port_desired_state(), port_name)
def reset(self):
"""
Sets up the environment according to the config
:return:
:rtype:
"""
context = self.get_mutable_context()
# self.simulator.SetDefaultContext(context)
# self.diagram.SetDefaultContext(context)
# put the ycb model slightly above the table
q_ycb_model = np.array([1, 0, 0, 0, 0, 0, 0.1])
self.set_slider_position(context, q_ycb_model)
self.set_slider_velocity(context, np.zeros(6))
# set pusher initial position
q_pusher = np.array([-0.1, 0.])
self.set_pusher_position(context, q_pusher)
self.set_pusher_velocity(context, np.zeros(2))
# set PID setpoint to zero
pid_setpoint = np.zeros(4)
self._pid_input_port_desired_state.FixValue(context, pid_setpoint)
self._sim_initialized = False
def set_initial_condition(self, q_slider, q_pusher):
"""
Sets the initial condition
:param q_slider:
:type q_slider:
:param q_pusher:
:type q_pusher:
:return:
:rtype:
"""
context = self.get_mutable_context()
self.reset() # sets velocities to zero
self.set_pusher_position(context, q_pusher)
self.set_slider_position(context, q_slider)
self.step(np.array([0, 0]), dt=2.0) # allow box to drop down
def set_initial_condition_from_dict(self, initial_cond):
self.set_initial_condition(
q_slider=initial_cond['q_slider'],
q_pusher=initial_cond['q_pusher'],
)
def step(self, action, dt=None):
assert action.size == 2
if not self._sim_initialized:
self._simulator.Initialize()
self._sim_initialized = True
if dt is None:
dt = self._step_dt
# the time to advance to
t_advance = self.get_context().get_time() + dt
context = self.get_mutable_context()
# set the value of the PID controller input port
pid_setpoint = np.concatenate((np.zeros(2), action))
self._pid_input_port_desired_state.FixValue(context, pid_setpoint)
# self._velocity_controller.set_velocity(action)
self._simulator.AdvanceTo(t_advance)
# set the sim time from the above context to avoid floating point errors
obs = self.get_observation()
# just placeholders for now
reward = 0.
done = False
info = None
return obs, reward, done, info
def render(self, mode='rgb'):
pass
def get_observation(
self,
context=None, # the specific context to use, otherwise none
):
"""
Gets obesrvations (optionally including images)
:param context:
:type context:
:return:
:rtype:
"""
if context is None:
context = self.get_context()
obs = dict()
obs["slider"] = self._dps_wrapper.get_model_state(
self.diagram, context, self._model_name)
obs['pusher'] = self._dps_wrapper.get_pusher_state(
self.diagram, context)
obs['sim_time'] = context.get_time()
if self._config["env"]["observation"]["rgbd"]:
obs["images"] = self._dps_wrapper.get_image_observations(
self.diagram, context)
return obs
def get_mutable_context(self):
return self._simulator.get_mutable_context()
def get_context(self):
return self._simulator.get_context()
def get_rgb_image(
self,
sensor_name: str,
context,
):
sensor = self._dps_wrapper._rgbd_sensors[sensor_name]
print("type(sensor)", type(sensor))
sensor_context = self.diagram.GetSubsystemContext(sensor, context)
rgb_img_PIL = drake_image_utils.get_color_image(sensor, sensor_context)
return rgb_img_PIL
def set_pusher_position(self, context, q):
"""
Sets the pusher position
:param context:
:type context:
:return:
:rtype:
"""
assert q.size == 2
mbp = self._dps_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetPositions(mbp_context, self._dps_wrapper.models['pusher'], q)
def set_pusher_velocity(self, context, v):
assert v.size == 2
mbp = self._dps_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetVelocities(mbp_context, self._dps_wrapper.models['pusher'], v)
def set_slider_position(self, context, q):
"""
Sets the slider position
:param context:
:type context:
:param q:
:type q:
:return:
:rtype:
"""
assert q.size == 7
# check that first 4 are a quaternion
if abs(np.linalg.norm(q[0:4]) - 1) > 1e-3:
print("q[0:4]", q[0:4])
print("np.linalg.norm(q[0:4])", np.linalg.norm(q[0:4]))
assert (abs(np.linalg.norm(q[0:4]) - 1) <
1e-3), "q[0:4] is not a valid quaternion"
# put the ycb model slightly above the table
mbp = self._dps_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetPositions(mbp_context,
self._dps_wrapper.models[self._model_name], q)
def set_slider_velocity(self, context, v):
assert v.size == 6
# put the ycb model slightly above the table
mbp = self._dps_wrapper.mbp
mbp_context = self._diagram.GetMutableSubsystemContext(mbp, context)
mbp.SetVelocities(mbp_context,
self._dps_wrapper.models[self._model_name], v)
def set_simulator_state_from_observation_dict(
self,
context,
d, # dictionary returned by DrakeSimEpisodeReader
):
# set slider position
q_slider = d['slider']['position']['raw']
v_slider = d['slider']['velocity']['raw']
self.set_slider_position(context, q_slider)
self.set_slider_velocity(context, v_slider)
# set pusher position
q_pusher = d['pusher']['position']
v_pusher = d['pusher']['velocity']
self.set_pusher_position(context, q_pusher)
self.set_pusher_velocity(context, v_pusher)
def pusher_within_boundary(self, context=None):
"""
Return true if pusher within boundary of table + tolerance
:return:
:rtype:
"""
if context is None:
context = self.get_context()
tol = 0.03 # 3 cm
pusher_state = self._dps_wrapper.get_pusher_state(
self.diagram, context)
pos = pusher_state['position'] # x,y
dims = np.array(self.config['env']['table']['size'])
high = dims[0:2] / 2.0 - tol
low = -high
in_range = ((low < pos) & (pos < high)).all()
return in_range
def slider_within_boundary(self, context=None):
"""
Return true if slider within boundary of table + tolerance
:return:
:rtype:
"""
if context is None:
context = self.get_context()
tol = 0.03 # 3 cm
slider_state = self._dps_wrapper.get_model_state(
self.diagram, context, self._model_name)
pos = slider_state['position']['translation']
dims = np.array(self.config['env']['table']['size'])
high = np.zeros(3)
low = np.zeros(3)
high[0:2] = dims[0:2] / 2.0 - tol
high[2] = 100 # just large number so that it's nonbinding
low[0:2] = -high[0:2]
low[2] = 0
in_range = ((low < pos) & (pos < high)).all()
return in_range
def state_is_valid(self): # checks if state is valid
return (self.slider_within_boundary()
and self.pusher_within_boundary())
def get_metadata(self):
"""
Returns a bunch of data that is useful to be logged
:return: dict()
:rtype:
"""
label_db = self._dps_wrapper.get_label_db()
mask_db = self._dps_wrapper.get_labels_to_mask()
# note saving TinyDB object with pickle causes error on load
# this is a workaround, can easily re-create TinyDB object
# when we load this
return {'label_db': label_db.all(), 'mask_db': mask_db.all()}
def get_labels_to_mask_list(self):
"""
Returns a list of labels that should be masked
:return:
:rtype:
"""
return self._dps_wrapper.get_labels_to_mask_list()
@property
def simulator(self):
return self._simulator
@property
def diagram(self):
return self._diagram
@property
def config(self):
return self._config
@property
def diagram_wrapper(self):
return self._dps_wrapper
def camera_pose(self, camera_name): # [4,4] homogeneous transform
pose_dict = self.config['env']['rgbd_sensors']['sensor_list'][
camera_name]
return transform_utils.transform_from_pose_dict(pose_dict)
def get_zero_action(self):
"""
Returns a zero action
:return:
:rtype:
"""
return np.zeros(2)
def camera_K_matrix(self, camera_name): # np.array [4,4]
return DrakePusherSliderEnv.camera_K_matrix_from_config(
self.config, camera_name)
@staticmethod
def camera_K_matrix_from_config(config, camera_name): # np.array [4,4]
sensor_dict = config['env']['rgbd_sensors']['sensor_list'][camera_name]
width = sensor_dict['width']
height = sensor_dict['height']
fov_y = sensor_dict['fov_y']
return transform_utils.camera_K_matrix_from_FOV(width, height, fov_y)
@staticmethod
def object_position_from_observation(obs):
q_slider = obs['slider']['position']['raw']
return env_utils.drake_position_to_pose(q_slider)
