forked from pz4kybsvg/Conception
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
70 lines
2.2 KiB
70 lines
2.2 KiB
"""A main() program (plus an reusable standalone function) that simulates a
|
|
spong-controlled acrobot.
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
|
|
from pydrake.examples import (
|
|
AcrobotPlant,
|
|
AcrobotSpongController,
|
|
AcrobotState,
|
|
)
|
|
from pydrake.systems.analysis import Simulator
|
|
from pydrake.systems.framework import DiagramBuilder
|
|
from pydrake.systems.primitives import LogVectorOutput
|
|
|
|
from drake.examples.acrobot.acrobot_io import load_scenario, save_output
|
|
|
|
|
|
def simulate(*, initial_state, controller_params, t_final, tape_period):
|
|
"""Simulates an Acrobot + Spong controller from the given initial state and
|
|
parameters until the given final time. Returns the state sampled at the
|
|
given tape_period.
|
|
"""
|
|
builder = DiagramBuilder()
|
|
plant = builder.AddSystem(AcrobotPlant())
|
|
controller = builder.AddSystem(AcrobotSpongController())
|
|
|
|
builder.Connect(plant.get_output_port(0), controller.get_input_port(0))
|
|
builder.Connect(controller.get_output_port(0), plant.get_input_port(0))
|
|
state_logger = LogVectorOutput(plant.get_output_port(0), builder,
|
|
tape_period)
|
|
|
|
diagram = builder.Build()
|
|
simulator = Simulator(diagram)
|
|
context = simulator.get_mutable_context()
|
|
plant_context = diagram.GetMutableSubsystemContext(plant, context)
|
|
controller_context = diagram.GetMutableSubsystemContext(
|
|
controller, context)
|
|
|
|
plant_context.SetContinuousState(initial_state)
|
|
controller_context.get_mutable_numeric_parameter(0).SetFromVector(
|
|
controller_params)
|
|
|
|
simulator.AdvanceTo(t_final)
|
|
|
|
x_tape = state_logger.FindLog(context).data()
|
|
return x_tape
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--scenario", metavar="*.yaml", required=True,
|
|
help="Scenario file to load (required).")
|
|
parser.add_argument(
|
|
"--output", metavar="*.yaml", required=True,
|
|
help="Output file to save (required).")
|
|
args = parser.parse_args()
|
|
scenario = load_scenario(filename=args.scenario)
|
|
x_tape = simulate(**scenario)
|
|
text = save_output(x_tape=x_tape)
|
|
with open(args.output, 'w') as handle:
|
|
handle.write(text)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|