Update examples

JeanElsner · Nov 17, 2023 · b7740b6 · b7740b6
1 parent 9bf4168
commit b7740b6
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 35 deletions.
diff --git a/examples/assets/two_arm.xml b/examples/assets/two_arm.xml
@@ -15,6 +15,7 @@
         friction='0.4'/>
     <light directional='true' diffuse='.7 .7 .7' pos='1 .1 2' dir='0 -.1 -2' specular='.3 .3 .3' castshadow='true'/>
     <body pos="0 0 1" euler="0 0 30">
+      <joint type="hinge" axis="0 0 1" damping="2000" />
       <geom type="box" size=".15 .2 .15" />
       <geom type="cylinder" size=".1 .3" />
       <geom type="sphere" size=".2" pos="0 0 .4" />

diff --git a/examples/gripper.py b/examples/gripper.py
@@ -13,7 +13,7 @@
 import numpy as np
 from dm_env import specs
 
-from dm_robotics.panda import env_builder
+from dm_robotics.panda import environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -51,9 +51,9 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   # The Panda model and environment support many customization
   # parameters. Here we use only the defaults and robot IP if provided.
   robot_params = params.RobotParams(robot_ip=args.robot_ip)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(robot_params)
+  panda_env = environment.PandaEnvironment(robot_params)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification.
     utils.full_spec(env)
     # Initialize the agent.

diff --git a/examples/haptics.py b/examples/haptics.py
@@ -8,10 +8,10 @@
 
 import dm_env
 import numpy as np
-from dm_control import mjcf
+from dm_control import composer, mjcf
 from dm_env import specs
 
-from dm_robotics.panda import arm_constants, env_builder
+from dm_robotics.panda import arm_constants, environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import utils
 
@@ -41,20 +41,16 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
 
   # Load environment from an MJCF file.
   XML_PATH = os.path.join(os.path.dirname(__file__), 'assets', 'haptics.xml')
-  arena = mjcf.from_file(XML_PATH)
+  arena = composer.Arena(xml_path=XML_PATH)
 
   # Robot parameters include the robot's IP for HIL,
   # haptic actuation mode and joint damping.
   robot_params = params.RobotParams(robot_ip=args.robot_ip,
                                     actuation=arm_constants.Actuation.HAPTIC,
                                     joint_damping=np.zeros(7))
-  # Environment parameters consist of the MJCF element defined above
-  # as arena and a higher control timestep for smooth interaction.
-  env_params = params.EnvirontmentParameters(arena=arena, control_timestep=0.01)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(robot_params,
-                                                          env_params=env_params)
+  panda_env = environment.PandaEnvironment(robot_params, arena, 0.01)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification.
     utils.full_spec(env)
     # Initialize the agent.

diff --git a/examples/minimal_working_example.py b/examples/minimal_working_example.py
@@ -3,7 +3,7 @@
 import numpy as np
 from dm_env import specs
 
-from dm_robotics.panda import env_builder
+from dm_robotics.panda import environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -33,9 +33,9 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   # The MoMa model and environment support many customization
   # parameters. Here we use only the defaults.
   robot_params = params.RobotParams(robot_ip=args.robot_ip)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(robot_params)
+  panda_env = environment.PandaEnvironment(robot_params)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification
     utils.full_spec(env)
     # Initialize the agent

diff --git a/examples/motion_cartesian.py b/examples/motion_cartesian.py
@@ -7,7 +7,7 @@
 import numpy as np
 from dm_env import specs
 
-from dm_robotics.panda import env_builder
+from dm_robotics.panda import environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -49,9 +49,9 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   # The Panda model and environment support many customization
   # parameters. Here we use only the defaults and robot IP if provided.
   robot_params = params.RobotParams(robot_ip=args.robot_ip)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(robot_params)
+  panda_env = environment.PandaEnvironment(robot_params)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification.
     utils.full_spec(env)
     # Initialize the agent.

diff --git a/examples/motion_joint.py b/examples/motion_joint.py
@@ -5,7 +5,7 @@
 import numpy as np
 from dm_env import specs
 
-from dm_robotics.panda import arm_constants, env_builder
+from dm_robotics.panda import arm_constants, environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -41,9 +41,9 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   # the actuation mode to joint velocities.
   robot_params = params.RobotParams(
       robot_ip=args.robot_ip, actuation=arm_constants.Actuation.JOINT_VELOCITY)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(robot_params)
+  panda_env = environment.PandaEnvironment(robot_params)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification.
     utils.full_spec(env)
     # Initialize the agent.

diff --git a/examples/multiple_robots.py b/examples/multiple_robots.py
@@ -3,7 +3,7 @@
 import numpy as np
 from dm_env import specs
 
-from dm_robotics.panda import env_builder
+from dm_robotics.panda import environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -38,10 +38,9 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
                                pose=[.5, -.5, 0, 0, 0, np.pi * 3 / 4])
   robot_3 = params.RobotParams(name='robot_3',
                                pose=[.5, .5, 0, 0, 0, np.pi * 5 / 4])
-  panda_env_builder = env_builder.PandaEnvironmentBuilder(
-      [robot_1, robot_2, robot_3])
+  panda_env = environment.PandaEnvironment([robot_1, robot_2, robot_3])
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification
     utils.full_spec(env)
     # Initialize the agent

diff --git a/examples/two_arm_robot.py b/examples/two_arm_robot.py
@@ -3,10 +3,10 @@
 
 import dm_env
 import numpy as np
-from dm_control import mjcf
+from dm_control import composer, mjcf
 from dm_env import specs
 
-from dm_robotics.panda import env_builder
+from dm_robotics.panda import environment
 from dm_robotics.panda import parameters as params
 from dm_robotics.panda import run_loop, utils
 
@@ -39,10 +39,10 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   # The environment includes a simple robot frame and
   # reference frame for our two-arm robot.
   XML_PATH = os.path.join(os.path.dirname(__file__), 'assets', 'two_arm.xml')
-  arena = mjcf.from_file(XML_PATH)
-  left_frame = arena.find('site', 'left')
-  right_frame = arena.find('site', 'right')
-  control_frame = arena.find('site', 'control')
+  arena = composer.Arena(xml_path=XML_PATH)
+  left_frame = arena.mjcf_model.find('site', 'left')
+  right_frame = arena.mjcf_model.find('site', 'right')
+  control_frame = arena.mjcf_model.find('site', 'control')
 
   # We use the sites defined in the MJCF to attach the robot arms
   # to the body. The robot's control site is used as a reference frame
@@ -53,11 +53,10 @@ def step(self, timestep: dm_env.TimeStep) -> np.ndarray:
   right = params.RobotParams(attach_site=right_frame,
                              name='right',
                              control_frame=control_frame)
-  env_params = params.EnvirontmentParameters(arena=arena)
-  panda_env_builder = env_builder.PandaEnvironmentBuilder([left, right],
-                                                          env_params)
+  env_params = params.EnvirontmentParameters(mjcf_root=arena)
+  panda_env = environment.PandaEnvironment([left, right], arena)
 
-  with panda_env_builder.build_task_environment() as env:
+  with panda_env.build_task_environment() as env:
     # Print the full action, observation and reward specification
     utils.full_spec(env)
     # Initialize the agent