Meshcat not showing the changes to a Free Body's Pose

Question

I've been trying to create my own ManipulationStation for a different robot arm using Pydrake, but I've been unsuccessful so far in adding clutter to my ManipulationStation. For some odd reason, Meshcat won't show the updated poses of my objects.

import numpy as np
import glob
from pydrake.geometry import MeshcatVisualizerCpp
from pydrake.math import RigidTransform, RotationMatrix
from pydrake.systems.analysis import Simulator

from pydrake.systems.framework import DiagramBuilder

from pydrake.all import (
    DiagramBuilder, FindResourceOrThrow,
    SceneGraph, Diagram,
    MultibodyPlant, Parser, Simulator, MeshcatVisualizerCpp, 
    UniformlyRandomRotationMatrix, RandomGenerator)

from pydrake.geometry import Meshcat 

class DexterPPStation(Diagram):
    def __init__(self, time_step, file_path):
        super().__init__()
        self.time_step = time_step
        self.path = file_path
        self.plant = MultibodyPlant(self.time_step)
        self.scene_graph = SceneGraph()
        self.plant.RegisterAsSourceForSceneGraph(self.scene_graph)
        self.controller_plant = MultibodyPlant(self.time_step)
        self.object_ids = []
        self.object_poses = []
        
    def AddObject(self, file, name, pose):
        model_idx = Parser(self.plant).AddModelFromFile(file, name)
        indices = self.plant.GetBodyIndices(model_idx)
        self.object_ids.append(indices[0])
        self.object_poses.append(pose)

        return model_idx

    def CreateBins(self, path, XP_B1, XP_B2):
        bin1 = Parser(self.plant).AddModelFromFile(path, "bin1")
        self.plant.WeldFrames(self.plant.world_frame(), self.plant.GetFrameByName("bin_base", bin1), XP_B1)
  
        bin2 = Parser(self.plant).AddModelFromFile(path, "bin2")
        self.plant.WeldFrames(self.plant.world_frame(), self.plant.GetFrameByName("bin_base", bin2), XP_B2)

    def CreateRandomPickingObjects(self, n = 4):
        choices = [f for f in glob.glob("/opt/drake/share/drake/manipulation/models/ycb/sdf/*.sdf")]
        z = 0.1
        rs = np.random.RandomState()
        generator = RandomGenerator(rs.randint(1000))
        for i in range(n):
            obj = choices[i]
            pose = RigidTransform(
                UniformlyRandomRotationMatrix(generator),  
                [rs.uniform(.35,0.6), rs.uniform(-.2, .2), z])
   
            model = self.AddObject(obj, obj.split("/")[-1].split(".")[0] + str(i), pose)
            body_idx = self.plant.GetBodyIndices(model)[0]
            self.object_ids.append(body_idx)
            self.object_poses.append(pose)
            z+=0.1
 
    def SetRandomPoses(self, station_context):
        plant_context = self.GetSubsystemContext(self.plant, station_context)
        for i in range(len(self.object_ids)):
            self.plant.SetFreeBodyPose(plant_context, self.plant.get_body(self.object_ids[i]), self.object_poses[i])
   
    def Finalize(self):
        self.plant.Finalize()
        self.controller_plant.Finalize()
        builder = DiagramBuilder()
        

        builder.AddSystem(self.plant)
        builder.AddSystem(self.controller_plant)
        builder.AddSystem(self.scene_graph)

        builder.Connect(self.plant.get_geometry_poses_output_port(), self.scene_graph.get_source_pose_port(self.plant.get_source_id()))
        builder.Connect(self.scene_graph.get_query_output_port(), self.plant.get_geometry_query_input_port())
        builder.ExportOutput(self.scene_graph.get_query_output_port(), "query_object")
        builder.ExportOutput(self.plant.get_geometry_poses_output_port(), "geometry_poses")

        builder.ExportOutput(self.scene_graph.get_query_output_port(), "geometry_query")

        builder.ExportOutput(self.plant.get_contact_results_output_port(),"contact_results")
        builder.ExportOutput(self.plant.get_state_output_port(),"plant_continuous_state")
        builder.BuildInto(self)

To test my code, I've been running the script below.

def test():
    builder = DiagramBuilder()
    station = DexterPPStation(1e-4, "drake/manipulation/models/final_dexter_description/urdf/dexter.urdf")

    station.CreateBins("/opt/drake/share/drake/examples/manipulation_station/models/bin.sdf", RigidTransform(np.array([0.5,0,0])), RigidTransform(np.array([0,0.5,0])))
    station.CreateRandomPickingObjects(1)
    
    station.Finalize()
    builder.AddSystem(station)
    station_context = station.CreateDefaultContext()
    station.SetRandomPoses(station_context)
    
    MeshcatVisualizerCpp.AddToBuilder(builder, station.GetOutputPort("query_object"), meshcat)

    diagram = builder.Build()
    simulator = Simulator(diagram)
    simulator.set_target_realtime_rate(1.0)
    simulator.AdvanceTo(0.1)

test()

I've tried to call the SetRandomPoses() function from inside my Finalize() method, but since I needed to pass in a context to the function, I wasn't sure what to do. I'm new to Drake, so any input would be greatly appreciated.

Answer 1

You've created a station_context and set it to the random poses, but then you don't use it anywhere. When you create the simulator, it is creating another Context (with the default values), which is getting published when you call AdvanceTo.

The solution here, I think, is to not create your own station_context, but do e.g.

simulator = Simulator(diagram)
diagram_context = simulator.get_mutable_context()
station_context = station.GetMyMutableContextFromRoot(diagram_context)
station.SetRandomPoses(station_context)

then you can call AdvanceTo.