Search code examples
drake

PyDrake: MonteCarloSimulation pose of freebody being reset


I've been struggling to debug this issue for a while but I can't figure out what's going on. I have a free body (Cube) that I set the default pose of before everything else (translation: (2, 1, 0)). When the make_simulator function (that's given to MonteCarloSimulation) runs, the pose of the free body is correctly set to the default. However, very mysteriously, when the output evaluation callback function runs, the pose.translation() of the free body has been reset to 0,0,0.

Appreciate your help in figuring out what I'm doing wrong/what's going on!

Here is a simple script that reproduces the issue:

import numpy as np
np.set_printoptions(suppress=True)
np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)})

from pydrake.all import (
    AddMultibodyPlantSceneGraph, Box, DiagramBuilder, MeshcatVisualizer, RigidTransform, RotationMatrix, 
    Simulator, SpatialInertia, Sphere, UnitInertia, 
    StartMeshcat, CoulombFriction, MonteCarloSimulation, RandomGenerator
)

# Start the visualizer.
meshcat = StartMeshcat()

def AddBox(plant, shape, name, mass=1, mu=1, color=[.5, .5, .9, 1.0]):
    instance = plant.AddModelInstance(name)
    inertia = UnitInertia.SolidBox(shape.width(), shape.depth(),
                                       shape.height())
     
    body = plant.AddRigidBody(
        name, instance,
        SpatialInertia(mass=mass,
                       p_PScm_E=np.array([0., 0., 0.]),
                       G_SP_E=inertia))
    if plant.geometry_source_is_registered():
        new_box_dims = np.array([x-(2 * 1e-7) for x in shape.size()])
        smaller_box = Box(*new_box_dims)

        sphere = Sphere(1e-7)
        sphere_transforms = [
            RigidTransform(RotationMatrix(),new_box_dims * 0.5),
            RigidTransform(RotationMatrix(),new_box_dims * -0.5),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [-1, 1, 1]),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [1, -1, 1]),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [1, 1, -1]),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [1, -1, -1]),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [-1, -1, 1]),
            RigidTransform(RotationMatrix(),new_box_dims * 0.5 * [-1, 1, -1]),
            ]
        for i, t in enumerate(sphere_transforms):
            sphere_name = f"point_{i}"
            plant.RegisterVisualGeometry(body, t, sphere, sphere_name, color)
            plant.RegisterCollisionGeometry(body, t, sphere, sphere_name,
                                        CoulombFriction(mu, mu))       
        
        plant.RegisterCollisionGeometry(body, RigidTransform(), smaller_box, name,
                                        CoulombFriction(mu, mu))        
        plant.RegisterVisualGeometry(body, RigidTransform(), shape, name, color)
    return body

def run_demo():
    mass = 1
    mu = 0.1
    mu_g = 0.1
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.01)
    box_geometry = Box(10,10,0.1) 
    X_WBox = RigidTransform(RotationMatrix(), [0, 0, -0.6])
    plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box_geometry, "ground",
                                    CoulombFriction(mu_g, mu_g))
    plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box_geometry, "ground",
                                    [.9, .9, .9, 1.0])
    
    box_geometry = Box(1,1,1) 
    
    X_WBox_Default = RigidTransform(RotationMatrix(), [2, 1, 0])
    
    box = AddBox(plant, box_geometry, "box", 
                        mass, mu, color=[0.8, 0, 0, 1.0])
    
    plant.SetDefaultFreeBodyPose(box, X_WBox_Default)

    plant.Finalize()

    MeshcatVisualizer.AddToBuilder(
        builder, scene_graph, meshcat)

             
    diagram = builder.Build()
    
    # Set up a simulator to run this diagram
    simulator = Simulator(diagram)
    
    context = simulator.get_mutable_context()
    plant_context = plant.GetMyMutableContextFromRoot(context)
    simulator.set_target_realtime_rate(1.0)
    simulator.Initialize()
    simulator.AdvanceTo(3) # Wait 3 seconds to verify initial positions is correct
    print("Initial position: ", plant.GetPositions(plant_context))

    # Perform the Monte Carlo simulation.
    def make_simulator(generator):
        ''' Create a simulator for the system
            using the given generator. '''
        simulator = Simulator(diagram)
        sim_context = simulator.get_mutable_context()
        sim_plant_context = plant.GetMyMutableContextFromRoot(sim_context)
        simulator.set_target_realtime_rate(1.0)
        simulator.Initialize()
        simulator.AdvanceTo(2.0)
        box_pos = plant.GetPositions(sim_plant_context)
        print(f"\nmake_simulator box pos: {box_pos}")

        return simulator

    
    def calc_loss(system, calc_loss_context):
        ''' Given a context from the end of the simulation,
            calculate a loss - distance between goal state
            and current state '''
 
        calc_loss_plant_context = plant.GetMyContextFromRoot(calc_loss_context)
        box_pos = plant.GetPositions(calc_loss_plant_context)
        
        print(f"calc_loss box pos: {box_pos}")
        return 1
    
    results = MonteCarloSimulation(
                make_simulator=make_simulator, output=calc_loss,
                final_time=5.0, num_samples=10, generator=RandomGenerator())
    

    meshcat.DeleteAddedControls()
    meshcat.Delete()
    return

run_demo()

And the output is the following:

INFO:drake:Meshcat listening for connections at http://localhost:7002 Initial position: [1.000 -0.000 0.000 0.000 2.000 1.000 -0.050]

make_simulator box pos: [1.000 -0.000 0.000 0.000 2.000 1.000 -0.050] calc_loss box pos: [1.000 -0.000 0.000 -0.000 0.000 0.000 -0.050]

make_simulator box pos: [1.000 -0.000 0.000 0.000 2.000 1.000 -0.050] calc_loss box pos: [1.000 -0.000 0.000 -0.000 0.000 0.000 -0.050]

make_simulator box pos: [1.000 -0.000 0.000 0.000 2.000 1.000 -0.050] calc_loss box pos: [1.000 -0.000 0.000 -0.000 0.000 0.000 -0.050] ...

EDIT: Additional Note: This issue doesn't arise if the box is not a free body. E.g. if I add a planar joint to the box then it doesn't reset it's pose when calc_loss is called.


Solution

  • The problem is that as of Drake v1.14.0, floating bodies have two different default poses -- the "default default" when used in CreateDefaultContext(), and the "random default" when used with SetRandomContext(). Calling SetDefaultFreeBodyPose() sets the former, but MonteCarlo uses the latter.

    This a bug in Drake. See https://github.com/RobotLocomotion/drake/issues/19106.

    In any case, you can work around it now by setting both defaults to what you want.

    You need to do this after calling Finalize. Here's the snippet:

        plant.Finalize()
        plant.SetDefaultFreeBodyPose(box, X_WBox_Default)
        # TODO SetDefaultFreeBodyPose is broken in Drake 1.14.0 when used in
        # MonteCarlo. Until that's fixed, we need to copy the default position into
        # the default "random" position as well.
        box_floating_joint = plant.GetJointByName("$world_box")
        box_floating_joint.set_random_position_distribution(
            X_WBox_Default.translation())
        box_floating_joint.set_random_quaternion_distribution(
            X_WBox_Default.rotation().ToQuaternion().cast[pydrake.symbolic.Expression]())