コード例 #1
0
ファイル: general_test.py プロジェクト: guofengw/drake
    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)
コード例 #2
0
ファイル: general_test.py プロジェクト: carismoses/drake
    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)
コード例 #3
0
ファイル: repro.py プロジェクト: EricCousineau-TRI/repro
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()))
コード例 #4
0
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()
コード例 #5
0
    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.)
コード例 #6
0
ファイル: lcm_test.py プロジェクト: thduynguyen/drake
 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()
コード例 #7
0
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.)
コード例 #9
0
    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)
コード例 #10
0
ファイル: joint_teleop.py プロジェクト: rpoyner-tri/drake
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)
コード例 #11
0
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
コード例 #12
0
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())
コード例 #13
0
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)
コード例 #14
0
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)