def test_explicit_steppers(self, explicit_stepper):
collective_system = ScalarExponentialDampedHarmonicOscillatorCollectiveSystem(
)
final_time = 1.0
n_steps = 500
stepper = explicit_stepper()
dt = np.float64(float(final_time) / n_steps)
time = np.float64(0.0)
tol = Tolerance.atol()
# Before stepping, let's extend the interface of the stepper
# while providing memory slots
from elastica.systems import make_memory_for_explicit_stepper
memory_collection = make_memory_for_explicit_stepper(
stepper, collective_system)
from elastica.timestepper import extend_stepper_interface
extend_stepper_interface(stepper, collective_system)
while np.abs(final_time - time) > 1e5 * tol:
time = stepper.do_step(collective_system, memory_collection, time,
dt)
for system in collective_system:
assert_allclose(
system.state,
system.analytical_solution(final_time),
rtol=Tolerance.rtol(),
atol=Tolerance.atol(),
)
def test_symplectic_stepper_throws_for_bad_stepper(self,
stepper_and_interface):
system = ScalarExponentialDecaySystem()
(stepper_cls, interface_cls) = stepper_and_interface
stepper = stepper_cls()
assert interface_cls not in stepper.__class__.__bases__
with pytest.raises(NotImplementedError) as excinfo:
extend_stepper_interface(stepper, system)
assert "steppers are supported" in str(excinfo.value)
def test_symplectic_stepper_interface_for_collective_systems(
self, stepper_and_interface):
system = SymplecticUndampedHarmonicOscillatorCollectiveSystem()
(stepper_cls, interface_cls) = stepper_and_interface
stepper = stepper_cls()
assert interface_cls not in stepper.__class__.__bases__
extend_stepper_interface(stepper, system)
assert interface_cls in stepper.__class__.__bases__
def test_symplectic_stepper_interface_for_simple_systems(
self, stepper_and_interface):
system = ScalarExponentialDecaySystem()
(stepper_cls, interface_cls) = stepper_and_interface
stepper = stepper_cls()
assert interface_cls not in stepper.__class__.__bases__
extend_stepper_interface(stepper, system)
assert interface_cls in stepper.__class__.__bases__
def test_symplectic_stepper_interface_for_simple_systems(
self, stepper_and_interface):
system = ScalarExponentialDecaySystem()
(stepper_cls, interface_cls) = stepper_and_interface
stepper = stepper_cls()
stepper_methods = None
assert stepper_methods is None
_, stepper_methods = extend_stepper_interface(stepper, system)
assert stepper_methods
def test_symplectic_stepper_interface_for_collective_systems(
self, stepper_and_interface):
system = SymplecticUndampedHarmonicOscillatorCollectiveSystem()
(stepper_cls, interface_cls) = stepper_and_interface
stepper = stepper_cls()
stepper_methods = None
assert stepper_methods is None
_, stepper_methods = extend_stepper_interface(stepper, system)
assert stepper_methods
def test_explicit_steppers(self, explicit_stepper):
collective_system = ScalarExponentialDampedHarmonicOscillatorCollectiveSystem(
)
final_time = 1.0
if explicit_stepper == EulerForward:
# Euler requires very small time-steps and in order not to slow down test,
# we are scaling the difference between analytical and numerical solution.
n_steps = 25000
scale = 1e3
else:
n_steps = 500
scale = 1
stepper = explicit_stepper()
dt = np.float64(float(final_time) / n_steps)
time = np.float64(0.0)
tol = Tolerance.atol()
# Before stepping, let's extend the interface of the stepper
# while providing memory slots
from elastica.systems import make_memory_for_explicit_stepper
memory_collection = make_memory_for_explicit_stepper(
stepper, collective_system)
from elastica.timestepper import extend_stepper_interface
do_step, stagets_and_updates = extend_stepper_interface(
stepper, collective_system)
while np.abs(final_time - time) > 1e5 * tol:
time = do_step(
stepper,
stagets_and_updates,
collective_system,
memory_collection,
time,
dt,
)
for system in collective_system:
assert_allclose(
system.state,
system.analytical_solution(final_time),
rtol=Tolerance.rtol() * scale,
atol=Tolerance.atol() * scale,
)
def reset(self):
"""
This class method, resets and creates the simulation environment. First,
Elastica rod (or arm) is initialized and boundary conditions acting on the rod defined.
Second, target and if there are obstacles are initialized and append to the
simulation. Finally, call back functions are set for Elastica rods and rigid bodies.
Returns
-------
"""
self.simulator = BaseSimulator()
# setting up test params
n_elem = self.n_elem
start = np.zeros((3, ))
start[1] = self.target_position[1]
direction = np.array([0.0, 0.0,
1.0]) # rod direction: pointing upwards
normal = np.array([0.0, 1.0, 0.0])
binormal = np.cross(direction, normal)
density = 1000
nu = self.NU # dissipation coefficient
E = self.E # Young's Modulus
poisson_ratio = 0.5
# Set the arm properties after defining rods
base_length = 1.0 # rod base length
radius_tip = 0.05 # radius of the arm at the tip
radius_base = 0.05 # radius of the arm at the base
radius_along_rod = np.linspace(radius_base, radius_tip, n_elem)
# Arm is shearable Cosserat rod
self.shearable_rod = CosseratRod.straight_rod(
n_elem,
start,
direction,
normal,
base_length,
base_radius=radius_along_rod,
density=density,
nu=nu,
youngs_modulus=E,
poisson_ratio=poisson_ratio,
)
# Now rod is ready for simulation, append rod to simulation
self.simulator.append(self.shearable_rod)
if self.mode != 2:
# fixed target position to reach
target_position = self.target_position
if self.mode == 2 or self.mode == 4:
# random target position to reach with boundary
t_x = np.random.uniform(self.boundary[0], self.boundary[1])
t_y = np.random.uniform(self.boundary[2], self.boundary[3])
if self.dim == 2.0 or self.dim == 2.5:
t_z = np.random.uniform(self.boundary[4], self.boundary[5]) * 0
elif self.dim == 3.0 or self.dim == 3.5:
t_z = np.random.uniform(self.boundary[4], self.boundary[5])
print("Target position:", t_x, t_y, t_z)
target_position = np.array([t_x, t_y, t_z])
# initialize target sphere
self.sphere = Sphere(
center=target_position, # initialize target position of the ball
base_radius=0.05,
density=1000,
)
if self.mode == 3:
self.dir_indicator = 1
self.sphere_initial_velocity = self.target_v
self.sphere.velocity_collection[..., 0] = [
self.sphere_initial_velocity,
0.0,
0.0,
]
if self.mode == 4:
self.rand_direction_1 = np.pi * np.random.uniform(0, 2)
if self.dim == 2.0 or self.dim == 2.5:
self.rand_direction_2 = np.pi / 2.0
elif self.dim == 3.0 or self.dim == 3.5:
self.rand_direction_2 = np.pi * np.random.uniform(0, 2)
self.v_x = (self.target_v * np.cos(self.rand_direction_1) *
np.sin(self.rand_direction_2))
self.v_y = (self.target_v * np.sin(self.rand_direction_1) *
np.sin(self.rand_direction_2))
self.v_z = self.target_v * np.cos(self.rand_direction_2)
self.sphere.velocity_collection[..., 0] = [
self.v_x,
self.v_y,
self.v_z,
]
self.boundaries = np.array(self.boundary)
# Set rod and sphere directors to each other.
self.sphere.director_collection[
..., 0] = self.shearable_rod.director_collection[..., 0]
self.simulator.append(self.sphere)
class WallBoundaryForSphere(FreeRod):
"""
This class generates a bounded space that sphere can move inside. If sphere
hits one of the boundaries (walls) of this space, it is reflected in opposite direction
with the same velocity magnitude.
"""
def __init__(self, boundaries):
self.x_boundary_low = boundaries[0]
self.x_boundary_high = boundaries[1]
self.y_boundary_low = boundaries[2]
self.y_boundary_high = boundaries[3]
self.z_boundary_low = boundaries[4]
self.z_boundary_high = boundaries[5]
def constrain_values(self, sphere, time):
pos_x = sphere.position_collection[0]
pos_y = sphere.position_collection[1]
pos_z = sphere.position_collection[2]
radius = sphere.radius
vx = sphere.velocity_collection[0]
vy = sphere.velocity_collection[1]
vz = sphere.velocity_collection[2]
if (pos_x - radius) < self.x_boundary_low:
sphere.velocity_collection[:] = np.array([-vx, vy, vz])
if (pos_x + radius) > self.x_boundary_high:
sphere.velocity_collection[:] = np.array([-vx, vy, vz])
if (pos_y - radius) < self.y_boundary_low:
sphere.velocity_collection[:] = np.array([vx, -vy, vz])
if (pos_y + radius) > self.y_boundary_high:
sphere.velocity_collection[:] = np.array([vx, -vy, vz])
if (pos_z - radius) < self.z_boundary_low:
sphere.velocity_collection[:] = np.array([vx, vy, -vz])
if (pos_z + radius) > self.z_boundary_high:
sphere.velocity_collection[:] = np.array([vx, vy, -vz])
def constrain_rates(self, sphere, time):
pass
if self.mode == 4:
self.simulator.constrain(self.sphere).using(
WallBoundaryForSphere, boundaries=self.boundaries)
# Add boundary constraints as fixing one end
self.simulator.constrain(self.shearable_rod).using(
OneEndFixedRod,
constrained_position_idx=(0, ),
constrained_director_idx=(0, ))
# Add muscle torques acting on the arm for actuation
# MuscleTorquesWithVaryingBetaSplines uses the control points selected by RL to
# generate torques along the arm.
self.torque_profile_list_for_muscle_in_normal_dir = defaultdict(list)
self.spline_points_func_array_normal_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_normal_dir,
muscle_torque_scale=self.alpha,
direction=str("normal"),
step_skip=self.step_skip,
torque_profile_recorder=self.
torque_profile_list_for_muscle_in_normal_dir,
)
self.torque_profile_list_for_muscle_in_binormal_dir = defaultdict(list)
self.spline_points_func_array_binormal_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_binormal_dir,
muscle_torque_scale=self.alpha,
direction=str("binormal"),
step_skip=self.step_skip,
torque_profile_recorder=self.
torque_profile_list_for_muscle_in_binormal_dir,
)
self.torque_profile_list_for_muscle_in_tangent_dir = defaultdict(list)
self.spline_points_func_array_tangent_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_tangent_dir,
muscle_torque_scale=self.beta,
direction=str("tangent"),
step_skip=self.step_skip,
torque_profile_recorder=self.
torque_profile_list_for_muscle_in_tangent_dir,
)
# Call back function to collect arm data from simulation
class ArmMuscleBasisCallBack(CallBackBaseClass):
"""
Call back function for Elastica rod
"""
def __init__(
self,
step_skip: int,
callback_params: dict,
):
CallBackBaseClass.__init__(self)
self.every = step_skip
self.callback_params = callback_params
def make_callback(self, system, time, current_step: int):
if current_step % self.every == 0:
self.callback_params["time"].append(time)
self.callback_params["step"].append(current_step)
self.callback_params["position"].append(
system.position_collection.copy())
self.callback_params["radius"].append(system.radius.copy())
self.callback_params["com"].append(
system.compute_position_center_of_mass())
self.callback_params["directors"].append(
system.director_collection.copy())
return
# Call back function to collect target sphere data from simulation
class RigidSphereCallBack(CallBackBaseClass):
"""
Call back function for target sphere
"""
def __init__(self, step_skip: int, callback_params: dict):
CallBackBaseClass.__init__(self)
self.every = step_skip
self.callback_params = callback_params
def make_callback(self, system, time, current_step: int):
if current_step % self.every == 0:
self.callback_params["time"].append(time)
self.callback_params["step"].append(current_step)
self.callback_params["position"].append(
system.position_collection.copy())
self.callback_params["radius"].append(
copy.deepcopy(system.radius))
self.callback_params["com"].append(
system.compute_position_center_of_mass())
self.callback_params["directors"].append(
system.director_collection.copy())
return
# Call back function to collect obstacle data from simulation
class RigidCylinderCallBack(CallBackBaseClass):
"""
Call back function for rigid cylinders or obstacles.
"""
def __init__(self, step_skip: int, callback_params: dict):
CallBackBaseClass.__init__(self)
self.every = step_skip
self.callback_params = callback_params
def make_callback(self, system, time, current_step: int):
if current_step % self.every == 0:
self.callback_params["time"].append(time)
self.callback_params["step"].append(current_step)
self.callback_params["position"].append(
system.position_collection.copy())
self.callback_params["com"].append(
system.compute_position_center_of_mass())
return
""" Add Obstacles to the environment """
self.obstacle = [None for _ in range(self.N_OBSTACLE)]
self.obstacle_histories = [
defaultdict(list) for _ in range(self.N_OBSTACLE)
]
self.obstacle_radii = [0.1 for _ in range(self.N_OBSTACLE)]
self.obstacle_length = [1.0 for _ in range(self.N_OBSTACLE)]
obstacle_direction = np.array([0.0, 1.0, 0.0])
self.obstacle_states = np.zeros(
(self.number_of_points_on_cylinder * self.N_OBSTACLE, 3))
if self.N_OBSTACLE > 0:
com_target = self.sphere.position_collection[..., -1]
start_position_array = np.zeros((3, self.N_OBSTACLE))
# first obstacle
start_position_array[
0, 0] = com_target[0] + 0.1 + self.obstacle_radii[0]
start_position_array[
1, 0] = com_target[1] - 0.5 * self.obstacle_length[0]
start_position_array[
2, 0] = com_target[2] - 0.3 * self.obstacle_radii[0]
# second obstacle
start_position_array[0, 1] = (com_target[0] + 0.1 +
3 * self.obstacle_radii[1] + 0.2)
start_position_array[
1, 1] = com_target[1] - 0.5 * self.obstacle_length[1]
start_position_array[
2, 1] = com_target[2] - 0.3 * self.obstacle_radii[1]
# third obstacle
start_position_array[0, 2] = (com_target[0] + 0.1 +
2 * self.obstacle_radii[2] + 0.1)
start_position_array[
1, 2] = com_target[1] - 0.5 * self.obstacle_length[2]
start_position_array[
2, 2] = com_target[2] + 2.2 * self.obstacle_radii[2]
# Forth obstacle underneath the third obstacle
start_position_array[0, 3] = start_position_array[0, 2]
start_position_array[1, 3] = start_position_array[1, 2]
start_position_array[2, 3] = (start_position_array[2, 1] -
0.2 * self.obstacle_radii[1] -
self.obstacle_radii[3])
# Fifth obstacle top of the third obstacle
start_position_array[0, 4] = start_position_array[0, 2]
start_position_array[1, 4] = start_position_array[1, 2]
start_position_array[2, 4] = (start_position_array[2, 3] -
self.obstacle_radii[1] -
self.obstacle_radii[4])
# Sixth obstacle top of the fifth obstacle
start_position_array[0, 5] = start_position_array[0, 4]
start_position_array[1, 5] = start_position_array[1, 4]
start_position_array[2, 5] = (start_position_array[2, 4] -
self.obstacle_radii[1] -
self.obstacle_radii[5])
# Seventh obstacle top of the third obstacle
start_position_array[0, 6] = start_position_array[0, 2]
start_position_array[1, 6] = start_position_array[1, 2]
start_position_array[2, 6] = (start_position_array[2, 2] +
self.obstacle_radii[1] +
self.obstacle_radii[6])
# Eight obstacle top of the third obstacle
start_position_array[0, 7] = start_position_array[0, 6]
start_position_array[1, 7] = start_position_array[1, 6]
start_position_array[2, 7] = (start_position_array[2, 6] +
self.obstacle_radii[6] +
self.obstacle_radii[7])
for i in range(self.N_OBSTACLE):
start = start_position_array[..., i]
self.obstacle[i] = Cylinder(
start=start,
direction=obstacle_direction,
normal=normal,
base_length=self.obstacle_length[i],
base_radius=self.obstacle_radii[i],
density=1000,
)
# Obstacle information for state
# obstacle_states(idx_on_obstacle, coordinates_x_y_z)
self.obstacle_states[
i * self.number_of_points_on_cylinder:(i + 1) *
self.number_of_points_on_cylinder, :, ] = (
start.reshape(3, 1) +
obstacle_direction.reshape(3, 1) *
np.linspace(0, self.obstacle_length[i],
self.number_of_points_on_cylinder)).T
self.obstacle_histories[i]["radius"] = self.obstacle_radii[i]
self.obstacle_histories[i]["height"] = self.obstacle_length[i]
self.obstacle_histories[i][
"direction"] = obstacle_direction.copy()
# We use scatter plot to plot obstacles, thus we need to discretize the obstacle and
# save the position of the discretized elements. This is only used for plotting and
# rendering purposes.
n_elem_for_plotting = 40
position_collection_for_plotting = np.zeros(
(3, n_elem_for_plotting))
end = (start_position_array[..., i] +
obstacle_direction * self.obstacle_length[i])
for k in range(0, 3):
position_collection_for_plotting[k, ...] = np.linspace(
start_position_array[..., i][k], end[k],
n_elem_for_plotting)
self.obstacle_histories[i][
"position_plotting"] = position_collection_for_plotting.copy(
)
self.simulator.append(self.obstacle[i])
# Constraint obstacle positions
self.simulator.constrain(self.obstacle[i]).using(
OneEndFixedRod,
constrained_position_idx=(0, ),
constrained_director_idx=(0, ),
)
# Add external contact
self.simulator.connect(
self.shearable_rod, self.obstacle[i]).using(
ExternalContact, k=8e4,
nu=4.0) # for rendering and plotting k=2*8e4, nu=4.0
""" Add Obstacles to the environment """
if self.COLLECT_DATA_FOR_POSTPROCESSING:
# Collect data using callback function for postprocessing
self.post_processing_dict_rod = defaultdict(list)
# list which collected data will be append
# set the diagnostics for rod and collect data
self.simulator.collect_diagnostics(self.shearable_rod).using(
ArmMuscleBasisCallBack,
step_skip=self.step_skip,
callback_params=self.post_processing_dict_rod,
)
self.post_processing_dict_sphere = defaultdict(list)
# list which collected data will be append
# set the diagnostics for cylinder and collect data
self.simulator.collect_diagnostics(self.sphere).using(
RigidSphereCallBack,
step_skip=self.step_skip,
callback_params=self.post_processing_dict_sphere,
)
for i in range(self.N_OBSTACLE):
self.simulator.collect_diagnostics(self.obstacle[i]).using(
RigidCylinderCallBack,
step_skip=self.step_skip,
callback_params=self.obstacle_histories[i],
)
# Finalize simulation environment. After finalize, you cannot add
# any forcing, constrain or call back functions
self.simulator.finalize()
# do_step, stages_and_updates will be used in step function
self.do_step, self.stages_and_updates = extend_stepper_interface(
self.StatefulStepper, self.simulator)
# set state
state = self.get_state()
# reset on_goal
self.on_goal = 0
# reset current_step
self.current_step = 0
# reset time_tracker
self.time_tracker = np.float64(0.0)
# reset previous_action
self.previous_action = None
# After resetting the environment return state information
return state
def reset(self):
"""
This class method, resets and creates the simulation environment. First,
Elastica rod (or arm) is initialized and boundary conditions acting on the rod defined.
Second, target and if there are obstacles are initialized and append to the
simulation. Finally, call back functions are set for Elastica rods and rigid bodies.
Returns
-------
"""
self.simulator = BaseSimulator()
# setting up test params
n_elem = self.n_elem
start = np.zeros((3,))
direction = np.array([0.0, 1.0, 0.0]) # rod direction: pointing upwards
normal = np.array([0.0, 0.0, 1.0])
binormal = np.cross(direction, normal)
density = 1000
nu = self.NU # dissipation coefficient
E = self.E # Young's Modulus
poisson_ratio = 0.5
# Set the arm properties after defining rods
base_length = 1.0 # rod base length
radius_tip = 0.05 # radius of the arm at the tip
radius_base = 0.05 # radius of the arm at the base
radius_along_rod = np.linspace(radius_base, radius_tip, n_elem)
# Arm is shearable Cosserat rod
self.shearable_rod = CosseratRod.straight_rod(
n_elem,
start,
direction,
normal,
base_length,
base_radius=radius_along_rod,
density=density,
nu=nu,
youngs_modulus=E,
poisson_ratio=poisson_ratio,
)
# Now rod is ready for simulation, append rod to simulation
self.simulator.append(self.shearable_rod)
# self.mode = 4
if self.mode != 2:
# fixed target position to reach
target_position = self.target_position
if self.mode == 2 or self.mode == 4:
# random target position to reach with boundary
t_x = np.random.uniform(self.boundary[0], self.boundary[1])
t_y = np.random.uniform(self.boundary[2], self.boundary[3])
if self.dim == 2.0 or self.dim == 2.5:
t_z = np.random.uniform(self.boundary[4], self.boundary[5]) * 0
elif self.dim == 3.0 or self.dim == 3.5:
t_z = np.random.uniform(self.boundary[4], self.boundary[5])
print("Target position:", t_x, t_y, t_z)
target_position = np.array([t_x, t_y, t_z])
# initialize sphere
self.sphere = Sphere(
center=target_position, # initialize target position of the ball
base_radius=0.05,
density=1000,
)
if self.mode == 3:
self.dir_indicator = 1
self.sphere_initial_velocity = self.target_v
self.sphere.velocity_collection[..., 0] = [
self.sphere_initial_velocity,
0.0,
0.0,
]
if self.mode == 4:
self.trajectory_iteration = 0 # for changing directions
self.rand_direction_1 = np.pi * np.random.uniform(0, 2)
if self.dim == 2.0 or self.dim == 2.5:
self.rand_direction_2 = np.pi / 2.0
elif self.dim == 3.0 or self.dim == 3.5:
self.rand_direction_2 = np.pi * np.random.uniform(0, 2)
self.v_x = (
self.target_v
* np.cos(self.rand_direction_1)
* np.sin(self.rand_direction_2)
)
self.v_y = (
self.target_v
* np.sin(self.rand_direction_1)
* np.sin(self.rand_direction_2)
)
self.v_z = self.target_v * np.cos(self.rand_direction_2)
self.sphere.velocity_collection[..., 0] = [
self.v_x,
self.v_y,
self.v_z,
]
self.boundaries = np.array(self.boundary)
# Set rod and sphere directors to each other.
self.sphere.director_collection[
..., 0
] = self.shearable_rod.director_collection[..., 0]
self.simulator.append(self.sphere)
class WallBoundaryForSphere(FreeRod):
"""
This class generates a bounded space that sphere can move inside. If sphere
hits one of the boundaries (walls) of this space, it is reflected in opposite direction
with the same velocity magnitude.
"""
def __init__(self, boundaries):
self.x_boundary_low = boundaries[0]
self.x_boundary_high = boundaries[1]
self.y_boundary_low = boundaries[2]
self.y_boundary_high = boundaries[3]
self.z_boundary_low = boundaries[4]
self.z_boundary_high = boundaries[5]
def constrain_values(self, sphere, time):
pos_x = sphere.position_collection[0]
pos_y = sphere.position_collection[1]
pos_z = sphere.position_collection[2]
radius = sphere.radius
vx = sphere.velocity_collection[0]
vy = sphere.velocity_collection[1]
vz = sphere.velocity_collection[2]
if (pos_x - radius) < self.x_boundary_low:
sphere.velocity_collection[:] = np.array([-vx, vy, vz])
if (pos_x + radius) > self.x_boundary_high:
sphere.velocity_collection[:] = np.array([-vx, vy, vz])
if (pos_y - radius) < self.y_boundary_low:
sphere.velocity_collection[:] = np.array([vx, -vy, vz])
if (pos_y + radius) > self.y_boundary_high:
sphere.velocity_collection[:] = np.array([vx, -vy, vz])
if (pos_z - radius) < self.z_boundary_low:
sphere.velocity_collection[:] = np.array([vx, vy, -vz])
if (pos_z + radius) > self.z_boundary_high:
sphere.velocity_collection[:] = np.array([vx, vy, -vz])
def constrain_rates(self, sphere, time):
pass
if self.mode == 4:
self.simulator.constrain(self.sphere).using(
WallBoundaryForSphere, boundaries=self.boundaries
)
# Add boundary constraints as fixing one end
self.simulator.constrain(self.shearable_rod).using(
OneEndFixedRod, constrained_position_idx=(0,), constrained_director_idx=(0,)
)
# Add muscle torques acting on the arm for actuation
# MuscleTorquesWithVaryingBetaSplines uses the control points selected by RL to
# generate torques along the arm.
self.torque_profile_list_for_muscle_in_normal_dir = defaultdict(list)
self.spline_points_func_array_normal_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_normal_dir,
muscle_torque_scale=self.alpha,
direction=str("normal"),
step_skip=self.step_skip,
max_rate_of_change_of_activation=self.max_rate_of_change_of_activation,
torque_profile_recorder=self.torque_profile_list_for_muscle_in_normal_dir,
)
self.torque_profile_list_for_muscle_in_binormal_dir = defaultdict(list)
self.spline_points_func_array_binormal_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_binormal_dir,
muscle_torque_scale=self.alpha,
direction=str("binormal"),
step_skip=self.step_skip,
max_rate_of_change_of_activation=self.max_rate_of_change_of_activation,
torque_profile_recorder=self.torque_profile_list_for_muscle_in_binormal_dir,
)
self.torque_profile_list_for_muscle_in_twist_dir = defaultdict(list)
self.spline_points_func_array_twist_dir = []
# Apply torques
self.simulator.add_forcing_to(self.shearable_rod).using(
MuscleTorquesWithVaryingBetaSplines,
base_length=base_length,
number_of_control_points=self.number_of_control_points,
points_func_array=self.spline_points_func_array_twist_dir,
muscle_torque_scale=self.beta,
direction=str("tangent"),
step_skip=self.step_skip,
max_rate_of_change_of_activation=self.max_rate_of_change_of_activation,
torque_profile_recorder=self.torque_profile_list_for_muscle_in_twist_dir,
)
# Call back function to collect arm data from simulation
class ArmMuscleBasisCallBack(CallBackBaseClass):
"""
Call back function for Elastica rod
"""
def __init__(
self, step_skip: int, callback_params: dict,
):
CallBackBaseClass.__init__(self)
self.every = step_skip
self.callback_params = callback_params
def make_callback(self, system, time, current_step: int):
if current_step % self.every == 0:
self.callback_params["time"].append(time)
self.callback_params["step"].append(current_step)
self.callback_params["position"].append(
system.position_collection.copy()
)
self.callback_params["radius"].append(system.radius.copy())
self.callback_params["com"].append(
system.compute_position_center_of_mass()
)
return
# Call back function to collect target sphere data from simulation
class RigidSphereCallBack(CallBackBaseClass):
"""
Call back function for target sphere
"""
def __init__(self, step_skip: int, callback_params: dict):
CallBackBaseClass.__init__(self)
self.every = step_skip
self.callback_params = callback_params
def make_callback(self, system, time, current_step: int):
if current_step % self.every == 0:
self.callback_params["time"].append(time)
self.callback_params["step"].append(current_step)
self.callback_params["position"].append(
system.position_collection.copy()
)
self.callback_params["radius"].append(copy.deepcopy(system.radius))
self.callback_params["com"].append(
system.compute_position_center_of_mass()
)
return
if self.COLLECT_DATA_FOR_POSTPROCESSING:
# Collect data using callback function for postprocessing
self.post_processing_dict_rod = defaultdict(list)
# list which collected data will be append
# set the diagnostics for rod and collect data
self.simulator.collect_diagnostics(self.shearable_rod).using(
ArmMuscleBasisCallBack,
step_skip=self.step_skip,
callback_params=self.post_processing_dict_rod,
)
self.post_processing_dict_sphere = defaultdict(list)
# list which collected data will be append
# set the diagnostics for cyclinder and collect data
self.simulator.collect_diagnostics(self.sphere).using(
RigidSphereCallBack,
step_skip=self.step_skip,
callback_params=self.post_processing_dict_sphere,
)
# Finalize simulation environment. After finalize, you cannot add
# any forcing, constrain or call back functions
self.simulator.finalize()
# do_step, stages_and_updates will be used in step function
self.do_step, self.stages_and_updates = extend_stepper_interface(
self.StatefulStepper, self.simulator
)
# set state
state = self.get_state()
# reset on_goal
self.on_goal = 0
# reset current_step
self.current_step = 0
# reset time_tracker
self.time_tracker = np.float64(0.0)
# reset previous_action
self.previous_action = None
# After resetting the environment return state information
return state
