def main():
    # Set up the world
    world = nimble.simulation.World()
    world.setGravity([0, -9.81, 0])
    world.setTimeStep(0.01)

    # Set up the box
    box = nimble.dynamics.Skeleton()
    boxJoint, boxBody = box.createTranslationalJoint2DAndBodyNodePair()
    boxShape = boxBody.createShapeNode(nimble.dynamics.BoxShape([.1, .1, .1]))
    boxVisual = boxShape.createVisualAspect()
    boxVisual.setColor([0.5, 0.5, 0.5])
    world.addSkeleton(box)

    # Set up initial conditions for optimization
    initial_position: torch.Tensor = torch.tensor([3.0, 0.0])
    initial_velocity: torch.Tensor = torch.zeros((world.getNumDofs()),
                                                 requires_grad=True)

    # Set up the GUI
    gui: nimble.NimbleGUI = nimble.NimbleGUI(world)
    gui.serve(8080)
    gui.nativeAPI().createSphere("goal",
                                 radius=0.1,
                                 pos=[0, 0, 0],
                                 color=[0, 255, 0])

    while True:
        state: torch.Tensor = torch.cat((initial_position, initial_velocity),
                                        0)
        states = [state]

        num_timesteps = 100
        for i in range(num_timesteps):
            state = nimble.timestep(world, state,
                                    torch.zeros((world.getNumDofs())))
            states.append(state)
        gui.loopStates(states)

        # Position is the first half of the state vector
        final_position = state[:world.getNumDofs()]

        # Our loss is just the distance to the origin at the final step
        loss = final_position.norm()
        print('loss: ' + str(loss))

        # Manually update weights using gradient descent. Wrap in torch.no_grad()
        # because weights have requires_grad=True, but we don't need to track this
        # in autograd.
        loss.backward()
        with torch.no_grad():
            learning_rate = 0.01
            initial_velocity -= learning_rate * initial_velocity.grad
            initial_velocity.grad = None

    gui.blockWhileServing()
Exemple #2
0
def main():
    world = nimble.simulation.World()
    world.setGravity([0, -9.81, 0])
    # Set up the 2D cartpole
    arm: nimble.dynamics.Skeleton = world.loadSkeleton(
        os.path.join(os.path.dirname(__file__),
                     "../../data/urdf/KR5/KR5 sixx R650.urdf"))
    ground: nimble.dynamics.Skeleton = world.loadSkeleton(
        os.path.join(os.path.dirname(__file__),
                     "../../data/sdf/atlas/ground.urdf"))
    floorBody: nimble.dynamics.BodyNode = ground.getBodyNode(0)
    floorBody.getShapeNode(0).getVisualAspect().setCastShadows(False)
    ticker = nimble.realtime.Ticker(world.getTimeStep())

    goal_x = 0.0
    goal_y = 0.8
    goal_z = -1.0
    goal: torch.Tensor = torch.tensor([goal_x, goal_y, goal_z])

    gui = nimble.NimbleGUI(world)
    gui.serve(8080)
    gui.nativeAPI().renderWorld(world, "world")
    gui.nativeAPI().createSphere("goal_pos", 0.1,
                                 np.array([goal_x, goal_y, goal_z]),
                                 np.array([0.0, 1.0, 0.0]), True, False)

    def onDrag(pos):
        nonlocal goal
        goal = torch.tensor(pos)
        gui.nativeAPI().setObjectPosition("goal_pos", pos)

    gui.nativeAPI().registerDragListener("goal_pos", onDrag)

    ikMap: nimble.neural.IKMapping = nimble.neural.IKMapping(world)
    handNode: nimble.dynamics.BodyNode = arm.getBodyNode("palm")
    ikMap.addLinearBodyNode(handNode)

    state: torch.Tensor = torch.randn((world.getStateSize()),
                                      requires_grad=True)

    learning_rate = 0.01

    while True:
        hand_pos: torch.Tensor = nimble.map_to_pos(world, ikMap, state)
        loss = (hand_pos - goal).square().sum()
        loss.backward()
        with torch.no_grad():
            state -= learning_rate * state.grad
            state.grad = None
        gui.nativeAPI().renderWorld(world, "world")
        time.sleep(0.002)

    gui.blockWhileServing()
# Set up the box
box = nimble.dynamics.Skeleton()
boxJoint, boxBody = box.createTranslationalJoint2DAndBodyNodePair()
boxShape = boxBody.createShapeNode(nimble.dynamics.BoxShape([.1, .1, .1]))
boxVisual = boxShape.createVisualAspect()
boxVisual.setColor([0.5, 0.5, 0.5])
world.addSkeleton(box)

# Set up initial conditions for optimization
initial_position: torch.Tensor = torch.tensor([3.0, 0.0])
initial_velocity: torch.Tensor = torch.zeros((world.getNumDofs()),
                                             requires_grad=True)

# Set up the GUI
gui: nimble.NimbleGUI = nimble.NimbleGUI(world)
gui.serve(8080)
gui.nativeAPI().createSphere("goal",
                             radius=0.1,
                             pos=[0, 0, 0],
                             color=[0, 255, 0])

while True:
    state: torch.Tensor = torch.cat((initial_position, initial_velocity), 0)
    states = [state]

    num_timesteps = 100
    for i in range(num_timesteps):
        state = nimble.timestep(world, state, torch.zeros(
            (world.getNumDofs())))
        states.append(state)
Exemple #4
0
import nimblephysics as nimble
import os


world: nimble.simulation.World = nimble.simulation.World()
arm: nimble.dynamics.Skeleton = world.loadSkeleton(os.path.join(
    os.path.dirname(__file__), "./KR5.urdf"))

# Your code here

gui = nimble.NimbleGUI(world)
gui.serve(8080)
gui.blockWhileServing()
Exemple #5
0
def main():
    # Load the world
    world: nimble.simulation.World = create_world()

    # Create cube
    projectile = create_projectile()
    world.addSkeleton(projectile)

    # Create catapult
    catapult = create_catapult()
    world.addSkeleton(catapult)

    floor = create_floor()
    world.addSkeleton(floor)

    # Target

    target_x = 2.2
    target_y = 2.2

    target = nimble.dynamics.Skeleton()
    target.setName('target')  # important for rendering shadows

    targetJoint, targetBody = floor.createWeldJointAndBodyNodePair()
    targetOffset = nimble.math.Isometry3()
    targetOffset.set_translation([target_x, target_y, 0])
    targetJoint.setTransformFromParentBodyNode(targetOffset)
    targetShape = targetBody.createShapeNode(
        nimble.dynamics.BoxShape([0.1, 0.1, 0.1]))
    targetVisual = targetShape.createVisualAspect()
    targetVisual.setColor([0.8, 0.5, 0.5])

    world.addSkeleton(target)

    gui: nimble.NimbleGUI = nimble.NimbleGUI(world)
    gui.serve(8080)

    # Set up the view

    # Define the loss function which is norm squared on final pose
    def loss(rollout: nimble.trajectory.TrajectoryRollout):
        poses = rollout.getPoses('identity')  # [n_bodies, steps]
        last_pos = poses[:, -1]
        last_x = last_pos[0]
        last_y = last_pos[1]
        final_loss = (target_x - last_x)**2 + (target_y - last_y)**2
        return final_loss

    dartLoss: nimble.trajectory.LossFn = nimble.trajectory.LossFn(loss)

    timesteps = 50
    shot_length = 20
    iteration_limit = 500
    trajectory = nimble.trajectory.MultiShot(world, dartLoss, timesteps,
                                             shot_length, False)
    trajectory.setParallelOperationsEnabled(True)

    # Initialize the optimizer
    optimizer = nimble.trajectory.IPOptOptimizer()
    optimizer.setLBFGSHistoryLength(5)
    optimizer.setTolerance(1e-4)
    optimizer.setCheckDerivatives(False)
    optimizer.setIterationLimit(iteration_limit)
    start = time.time()
    result: nimble.trajectory.Solution = optimizer.optimize(trajectory)
    print(f'Finished. Took: {time.time() - start}')

    # Get the rollout from the last optimization step.
    n_registered_opt_steps = result.getNumSteps()
    rollout: nimble.trajectory.TrajectoryRollout = result.getStep(
        n_registered_opt_steps - 1).rollout
    poses = rollout.getPoses()  # Get poses
    vels = rollout.getVels()  # Get velocities
    states = torch.cat((torch.tensor(poses), torch.tensor(vels)),
                       0).transpose(1, 0)

    dofs = world.getNumDofs()
    poses = np.zeros((dofs, len(states)))
    for i in range(len(states)):
        # Take the top-half of each state vector, since this is the position component
        poses[:, i] = states[i].detach().numpy()[:dofs]

    gui.loopStates(states)  # tells the GUI to animate our list of states
    gui.blockWhileServing()  # block here so we don't exit the program
Exemple #6
0
def main():
    world = nimble.simulation.World()
    world.setGravity([0, -9.81, 0])

    # Set up the 2D cartpole

    cartpole = nimble.dynamics.Skeleton()
    cartRail, cart = cartpole.createPrismaticJointAndBodyNodePair()
    cartRail.setAxis([1, 0, 0])
    cartShape = cart.createShapeNode(nimble.dynamics.BoxShape([.5, .1, .1]))
    cartVisual = cartShape.createVisualAspect()
    cartVisual.setColor([0.5, 0.5, 0.5])
    cartRail.setPositionUpperLimit(0, 10)
    cartRail.setPositionLowerLimit(0, -10)
    cartRail.setControlForceUpperLimit(0, 10)
    cartRail.setControlForceLowerLimit(0, -10)

    poleJoint, pole = cartpole.createRevoluteJointAndBodyNodePair(cart)
    poleJoint.setAxis([0, 0, 1])
    poleShape = pole.createShapeNode(nimble.dynamics.BoxShape([.1, 1.0, .1]))
    poleVisual = poleShape.createVisualAspect()
    poleVisual.setColor([0.7, 0.7, 0.7])
    poleJoint.setControlForceUpperLimit(0, 0)
    poleJoint.setControlForceLowerLimit(0, 0)

    poleOffset = nimble.math.Isometry3()
    poleOffset.set_translation([0, -0.5, 0])
    poleJoint.setTransformFromChildBodyNode(poleOffset)

    world.addSkeleton(cartpole)

    # Make simulations repeatable
    random.seed(1234)

    # Make simulations and backprop run faster by using a bigger timestep
    world.setTimeStep(world.getTimeStep() * 10)

    num_timesteps = 100

    action_size = world.getActionSize()
    print('action size: ' + str(action_size))
    learning_rate = 0.01

    first_state: torch.Tensor = torch.randn((world.getStateSize()),
                                            requires_grad=True)
    actions: List[torch.Tensor] = [
        torch.zeros((action_size), requires_grad=True)
        for _ in range(num_timesteps)
    ]

    gui: nimble.NimbleGUI = nimble.NimbleGUI(world)
    gui.serve(8080)

    while True:
        state: torch.Tensor = first_state
        states = [state]
        for i in range(num_timesteps):
            state = nimble.timestep(world, state, actions[i])
            states.append(state)
        gui.loopStates(states)

        loss = state.norm()
        print('loss: ' + str(loss))

        # Manually update weights using gradient descent. Wrap in torch.no_grad()
        # because weights have requires_grad=True, but we don't need to track this
        # in autograd.
        loss.backward()
        with torch.no_grad():
            first_state -= learning_rate * first_state.grad
            first_state.grad = None

            for action in actions:
                action -= learning_rate * action.grad
                action.grad = None
        # time.sleep(0.1)

    gui.blockWhileServing()